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