【leetcode】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.

解题思路：本题我用了三个list，一个是stack_list，保存所以的stack信息；一个是nonEmptyStack，记录当前不为空的stack的下标；还一个是availableStack，记录当前未满的stack的下标。在push的时候，只需要找出availableStack中所有下标的最小值，插入对应的stack即可，如果availableStack为空，新增一个stack，插入stack_list，同时更新availableStack和nonEmptyStack的状态；pop操作则是找出nonEmptyStack中下标的最大值，对其对应的stack做pop操作，而popAsStack的操作就更简单，可以直接用下标访问stack_list。需要注意的是，每次对stack有任何操作，都要同步更新availableStack和nonEmptyStack。因为是求availableStack和nonEmptyStack的最大或者最小值，只需要保证availableStack和nonEmptyStack中的元素有序即可，更新availableStack和nonEmptyStack则可以用二分查找法。

代码如下：

class DinnerPlates(object):

    def __init__(self, capacity):
        """
        :type capacity: int
        """
        self.stack_list = []
        self.availableStack = []
        self.capacity = capacity
        self.nonEmptyStack = []

    def push(self, val):
        """
        :type val: int
        :rtype: None
        """
        if len(self.availableStack) == 0:
            inx = len(self.stack_list)
            #self.availableStack.append(len(self.stack_list))
            self.stack_list.append([val])
            if len(self.stack_list[inx]) < self.capacity:
                self.availableStack.append(inx)
        else:
            inx = self.availableStack[0]
            self.stack_list[inx].append(val)
            if len(self.stack_list[inx]) >= self.capacity:
                self.availableStack.pop(0)
        import bisect
        b_inx = bisect.bisect_left(self.nonEmptyStack,inx)
        if b_inx == -1 or b_inx == len(self.nonEmptyStack) or self.nonEmptyStack[b_inx] != inx:
            bisect.insort_left(self.nonEmptyStack,inx)


    def pop(self):
        """
        :rtype: int
        """
        if len(self.nonEmptyStack) == 0:
            return -1
        inx = self.nonEmptyStack[-1]
        v = self.stack_list[inx].pop(-1)
        if len(self.stack_list[inx]) == 0:
            self.nonEmptyStack.pop(-1)
        return v


    def popAtStack(self, index):
        """
        :type index: int
        :rtype: int
        """
        import bisect
        b_inx = bisect.bisect_left(self.nonEmptyStack, index)
        if b_inx == -1 or b_inx == len(self.nonEmptyStack) or self.nonEmptyStack[b_inx] != index:
            return -1
        v = self.stack_list[index].pop(-1)
        if len(self.stack_list[index]) == 0:
            del self.nonEmptyStack[b_inx]

        b_inx = bisect.bisect_left(self.availableStack, index)
        if b_inx == -1 or b_inx == len(self.availableStack) or self.availableStack[b_inx] != index:
            bisect.insort_left(self.availableStack,index)

        return v