Introduction

In this blog post, we will explore how to solve the LeetCode problem "1172" using AI. LeetCode is a popular platform for preparing for coding interviews, and with the help of AI tools like Chatmagic, we can generate solutions quickly and efficiently - helping you pass the interviews and get the job offer without having to study for months.

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 capacity. void push(int val) Pushes the given integer val into the leftmost stack with a 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 the stacks are empty. int popAtStack(int index) Returns the value at the top of the stack with the given index index and removes it from that stack or returns -1 if the stack with that given index is empty. Example 1: Input ["DinnerPlates", "push", "push", "push", "push", "push", "popAtStack", "push", "push", "popAtStack", "popAtStack", "pop", "pop", "pop", "pop", "pop"] [[2], [1], [2], [3], [4], [5], [0], [20], [21], [0], [2], [], [], [], [], []] Output [null, null, null, null, null, null, 2, null, null, 20, 21, 5, 4, 3, 1, -1] Explanation: DinnerPlates D = DinnerPlates(2); // Initialize with capacity = 2 D.push(1); D.push(2); D.push(3); D.push(4); D.push(5); // The stacks are now: 2 4 1 3 5 ﹈ ﹈ ﹈ D.popAtStack(0); // Returns 2. The stacks are now: 4 1 3 5 ﹈ ﹈ ﹈ D.push(20); // The stacks are now: 20 4 1 3 5 ﹈ ﹈ ﹈ D.push(21); // The stacks are now: 20 4 21 1 3 5 ﹈ ﹈ ﹈ D.popAtStack(0); // Returns 20. The stacks are now: 4 21 1 3 5 ﹈ ﹈ ﹈ D.popAtStack(2); // Returns 21. The stacks are now: 4 1 3 5 ﹈ ﹈ ﹈ D.pop() // Returns 5. The stacks are now: 4 1 3 ﹈ ﹈ D.pop() // Returns 4. The stacks are now: 1 3 ﹈ ﹈ D.pop() // Returns 3. The stacks are now: 1 ﹈ D.pop() // Returns 1. There are no stacks. D.pop() // Returns -1. There are still no stacks. Constraints: 1 <= capacity <= 2 104 1 <= val <= 2 104 0 <= index <= 105 At most 2 * 105 calls will be made to push, pop, and popAtStack.

Explanation

• Maintain a list of stacks. Use two heaps (min-heap and max-heap) to efficiently track the available stacks for pushing (stacks that have space) and non-empty stacks for popping.
  • The min-heap stores indices of stacks that are not full. The max-heap stores indices of stacks that are not empty.
  • Lazy removal: When popAtStack is called, instead of physically removing the stack if it becomes empty, simply mark it as removed. Before popping from the rightmost stack, clean the max-heap to remove any indices that point to empty stacks.

• Runtime Complexity: O(log N) for push, pop, and popAtStack operations, where N is the number of stacks. Storage Complexity: O(N), where N is the number of stacks.

Code

    import heapq

class DinnerPlates:

    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = []
        self.available = []  # min heap for available stacks
        self.non_empty = []  # max heap for non-empty stacks
        self.size = 0

    def push(self, val: int) -> None:
        while self.available and self.available[0] not in range(len(self.stacks)):
            heapq.heappop(self.available)

        if self.available:
            index = heapq.heappop(self.available)
            self.stacks[index].append(val)
            heapq.heappush(self.non_empty, -index)
        else:
            self.stacks.append([val])
            index = len(self.stacks) - 1
            heapq.heappush(self.non_empty, -index)

        if len(self.stacks[-1]) < self.capacity:
            heapq.heappush(self.available, len(self.stacks) - 1)
        self.size = len(self.stacks)

    def pop(self) -> int:
        while self.non_empty and -self.non_empty[0] >= len(self.stacks):
                heapq.heappop(self.non_empty)

        while self.non_empty and not self.stacks[-self.non_empty[0]]:
                heapq.heappop(self.non_empty)

        if not self.non_empty:
            return -1

        index = -heapq.heappop(self.non_empty)
        val = self.stacks[index].pop()

        if len(self.stacks[index]) == 0:
            while self.available and self.available[0] not in range(len(self.stacks)):
                heapq.heappop(self.available)

        if index < len(self.stacks) - 1 and len(self.stacks[index]) < self.capacity:
            heapq.heappush(self.available, index)


        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()
            if self.available and self.available[-1] == len(self.stacks):
                self.available.pop()
        self.size = len(self.stacks)    
        return val

    def popAtStack(self, index: int) -> int:
        if index >= len(self.stacks) or not self.stacks[index]:
            return -1

        val = self.stacks[index].pop()
        if len(self.stacks[index]) == 0:
            while self.available and self.available[0] not in range(len(self.stacks)):
                heapq.heappop(self.available)

        if len(self.stacks[index]) < self.capacity:
            heapq.heappush(self.available, index)
        while self.non_empty and -self.non_empty[0] >= len(self.stacks):
            heapq.heappop(self.non_empty)
        while self.non_empty and not self.stacks[-self.non_empty[0]]:
                heapq.heappop(self.non_empty)

        while self.stacks and not self.stacks[-1]:
            self.stacks.pop()
            if self.available and self.available[-1] == len(self.stacks):
                self.available.pop()

        self.size = len(self.stacks)

        return val