-
Notifications
You must be signed in to change notification settings - Fork 891
/
problem_141.py
46 lines (38 loc) · 990 Bytes
/
problem_141.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class Stack:
def __init__(self):
self.list = []
def __repr__(self):
return str(self.list)
def pop(self, stack_number):
if not stack_number < len(self.list):
return
stack = self.list[stack_number]
if not stack:
return
val = stack.pop()
return val
def push(self, item, stack_number):
if stack_number < len(self.list):
stack = self.list[stack_number]
stack.append(item)
elif stack_number == len(self.list):
new_stack = list()
new_stack.append(item)
self.list.append(new_stack)
# Tests
st = Stack()
st.push(1, 0)
assert st.list == [[1]]
st.push(2, 1)
st.push(3, 2)
assert st.list == [[1], [2], [3]]
val = st.pop(3)
assert not val
val = st.pop(2)
assert val == 3
assert st.list == [[1], [2], []]
st.push(6, 0)
st.push(7, 0)
assert st.list == [[1, 6, 7], [2], []]
val = st.pop(0)
assert st.list == [[1, 6], [2], []]