AlgoDesign

Stack of plates

1class Solution:
2    def __init__(self, capacity):
3        self.stk = []
4        self.capacity = capacity
5
6    def push(self, val):
7        if len(self.stk) == 0 or len(self.stk[-1]) == self.capacity:
8            self.stk.append([val])
9            return
10        self.stk[-1].append(val)
11
12    def pop(self):
13        if len(self.stk) == 0: return None
14        res = self.stk[-1].pop()
15        if len(self.stk[-1]) == 0: self.stk.pop()
16        return res
17
18    def popAt(self, index):
19        if len(self.stk) == 0: return None
20        return self.leftShift(index, True)
21    
22    # rollover the bottom values from the next stacks onto the current stack's top at index where the top was popped
23    def leftShift(self, index, removeTop):
24        stack = self.stk[index]
25        removedItem = None
26        if removeTop: removedItem = stack.pop()
27        else: removedItem = self.removeBottom(index)
28
29        if len(self.stk[index]) == 0: del self.stk[index]
30        else:
31            val = self.leftShift(index + 1, False)
32            self.stk[index].append(val)
33        
34        return removedItem
35    
36    def removeBottom(self, index):
37        return self.stk[index].pop(0)