LeetCode- 1172. Dinner Plate Stacks

Problem statement :-

You have an infinite number of stacks arranged in a row and numbered (left to right) from 0, each of the stacks has the same maximum capacity.

Implement the DinnerPlates class:

DinnerPlates(int capacity) Initializes the object with the maximum capacity of the stacks. void push(int val) pushes the given positive integer val into the leftmost stack with size less than capacity. int pop() returns the value at the top of the rightmost non-empty stack and removes it from that stack, and returns -1 if all stacks are empty. int popAtStack(int index) returns the value at the top of the stack with the given index and removes it from that stack, and returns -1 if the stack with that given index is empty.

Example

D = DinnerPlates(2)
D.push(1)
D.push(2)
D.push(3)
D.push(4)
D.push(5)

# Stacks = [[1,2], [3,4],[5]]

op = D.popAtStack(0)

# op = 2
# stacks = [[1], [3,4],[5]]

D.push(20)

# stacks = [[1,20], [3,4],[5]]

D.push(21)

# stacks = [[1,20], [3,4],[5,21]]

D.pop()

# stacks = [[1,20], [3,4],[5]]

D.pop()

# stacks = [[1,20], [3,4]]

Solution

In this implementation of stacks we have to hack into fundamentals of how we normally push and pop from the stack. As we have popAtStack function which provides the index from the element is to be poped that leads to empty positions between the ends of stacks. So here is how we implement such class.

Before pushing the new element we have to check whether there is any empty space in the middle of the stack if yes then append to that particular list. If not then we have to check whether capacity is full or not if not append to the last stack of the stacks. If the capacity is reached add a new empty stack and append the element to it.

To keep track of empty positions caused by removing elements using popAtStack function we are using heap.

As mentioned in the problem statement, that if there is no elements left return -1 so before popping out we have to check does stack contains any other stacks, if yes then we have to remove the element from last stack, if not remove the empty stack.

Here is the solution.

import heapq
class DinnerPlates:
    def __init__(self, capacity):
        self.stack = []
        self.que = []
        self.capacity = capacity
    def push(self, val):
        i = self._get()
        if i != -1:
            self.stack[i].append(val)
        else:
            if not self.stack or len(self.stack[-1]) == self.capacity:
                self.stack.append([])
            self.stack[-1].append(val)
    def pop(self):
        if self.stack:
            if self.stack[-1]:
                v = self.stack[-1].pop()
                if self.stack and not self.stack[-1]:
                    self.stack.pop()
                return v
        return -1
    def popAtStack(self, index):
        if index >= len(self.stack) or not self.stack[index]:
            return -1
        v = self.stack[index].pop()
        heapq.heappush(self.que, index)
        return v

    def _get(self):
        if self.que:
            i = heapq.heappop(self.que)
            if i < len(self.stack) and len(stack[i]) < self.capacity:
                return i
        return -1