Problem statement

https://binarysearch.com/problems/Stack-of-Stacks/

Solution

Equal to Leetcode 1172. Dinner Plate Stacks.

Complexity

Time complexity for push is O(log n) for other two is O(1).

Code

class StackOfStacks:
    def __init__(self, capacity):
        self.plates = []
        self.c = capacity
        self.emp = []

    def push(self, val):
        if not self.emp:
            self.plates += [[val]]
            for _ in range(self.c - 1): heappush(self.emp, len(self.plates)-1)
            return
            
        ind = heappop(self.emp)
        if ind == len(self.plates): self.plates += [[]]
        self.plates[ind] += [val]
       

    def pop(self):
        while len(self.plates) > 0 and not self.plates[-1]:
            self.plates.pop()
            
        if not self.plates: return -1
        heapq.heappush(self.emp, len(self.plates) - 1)
        return self.plates[-1].pop()

    def popStack(self, ind):
        if ind >= len(self.plates) or not self.plates[ind]: return -1
        heapq.heappush(self.emp, ind)
        return self.plates[ind].pop()