Dinner Plate Stacks LeetCode solution in Python

1
class DinnerPlates:
2

3
    def __init__(self, capacity: int):
4
        self.heap = []
5
        self.stacks = []
6
        self.capacity = capacity
7

8
    def push(self, val: int) -> None:
9
        if self.heap:
10
            index = heapq.heappop(self.heap)
11
            if index < len(self.stacks):
12
                self.stacks[index].append(val)
13
            else:
14
                self.push(val)
15
        elif self.stacks:
16
            lastStack = self.stacks[-1]
17
            if len(lastStack) != self.capacity:
18
                lastStack.append(val)
19
            else:
20
                stack = deque()
21
                stack.append(val)
22
                self.stacks.append(stack)
23
        else:
24
            stack = deque()
25
            stack.append(val)
26
            self.stacks.append(stack)
27

28
    def pop(self) -> int:
29
        while self.stacks:
30
            lastStack = self.stacks[-1]
31
            if lastStack:
32
                val = lastStack.pop()
33
                if not lastStack:
34
                    self.stacks.pop()
35
                return val
36
            else:
37
                self.stacks.pop()
38
        return -1
39

40
    def popAtStack(self, index: int) -> int:
41
        if index == len(self.stacks) - 1:
42
            return self.pop()
43
        if index < len(self.stacks):
44
            stack = self.stacks[index]
45
            if stack:
46
                val = stack.pop()
47
                heapq.heappush(self.heap, index)
48
                return val
49
        return -1
50

51

52
# Your DinnerPlates object will be instantiated and called as such:
53
# obj = DinnerPlates(capacity)
54
# obj.push(val)
55
# param_2 = obj.pop()
56
# param_3 = obj.popAtStack(index)