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 maximumcapacity
of the stacks.void push(int val)
pushes the given positive integerval
into the leftmost stack with size less thancapacity
.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 givenindex
and removes it from that stack, and returns -1 if the stack with that givenindex
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 topush
,pop
, andpopAtStack
.
-----------------------------------------------------------------------------------
这题用堆并不难想到,难点是如何简化codes,有下面几个问题需要考虑清楚:
1. 堆里放什么?所有删过的栈下标都放进堆里。可能存在stack[stack]
2. 利用push_at和pop_at来简化代码
3. size,idx这些成员变量没有必要记,python的len复杂度是O(1)
import heapq
from collections import defaultdict
class DinnerPlates:
def __init__(self, capacity: int):
self.cap = capacity
self.pq,self.sta=[],[]
def push(self, val: int) -> None:
def push_at(index, val):
if (index < 0 or index >= len(self.sta) or len(self.sta[index]) == self.cap):
self.sta.append([val])
else:
self.sta[index].append(val)
insert_idx = heapq.heappop(self.pq) if self.pq else len(self.sta)-1
push_at(insert_idx, val)
def pop(self) -> int:
while (self.sta and (not self.sta[-1])):
self.sta.pop()
return self.popAtStack(len(self.sta)-1)
def popAtStack(self, index: int) -> int:
if (0 <= index < len(self.sta) and self.sta[index]):
heapq.heappush(self.pq,index)
return self.sta[index].pop()
return -1