1172. Dinner Plate Stacks

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:

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 <= 20000
• 1 <= val <= 20000
• 0 <= index <= 100000
• At most 200000 calls will be made to push, pop, and popAtStack.

思路：维护要pop和push的index，然后每次更新后都update对应的index

import heapq
class DinnerPlates(object):

    def __init__(self, capacity):
        """
        :type capacity: int
        """
        self.stacks = []
        self.capacity = capacity
        self.pop_pq = []
        self.push_pq = []

    def push(self, val):
        """
        :type val: int
        :rtype: None
        """
        while self.push_pq and len(self.stacks[self.push_pq[0]])>=self.capacity:
            heapq.heappop(self.push_pq)
        if len(self.push_pq)==0:
            next_used = len(self.stacks)
            heapq.heappush(self.push_pq, next_used)
            heapq.heappush(self.pop_pq, -next_used)
            self.stacks.append([val])
        else:
            push_idx = self.push_pq[0]
            self.stacks[push_idx].append(val)

    def pop(self):
        """
        :rtype: int
        """
        while self.pop_pq and len(self.stacks[-self.pop_pq[0]])==0:
            heapq.heappop(self.pop_pq)
        if len(self.pop_pq)==0:
            return -1
        else:
            pop_idx = -self.pop_pq[0]
            return self.popAtStack(pop_idx)
            
        

    def popAtStack(self, index):
        """
        :type index: int
        :rtype: int
        """
        if len(self.stacks[index])==0: return -1
        if len(self.stacks[index]) == self.capacity:
            heapq.heappush(self.push_pq, index)
        return self.stacks[index].pop()