题目描述
思路
这题目叫做无限栈来得更好一点。 就纯模拟：
-
push的时候从左向右看第一个未满的，push进去；
-
pop的时候，从右向左看，第一个有元素的，pop出来
代码
class DinnerPlates:
def __init__(self, capacity: int):
self.capacity = capacity
self.not_full_index = SortedSet() # 存放未满栈的下标（可能有空栈在其中）
self.top_index = 0 # 指向右侧第一个空栈(随着push增大，随着pop而缩小)
self.infinity_stk = [] # 模拟无限栈
def push(self, val: int) -> None: # 从左向右查看第一个未满的栈
if len(self.not_full_index) == 0 : # 说明添加刚刚开始或者目前全满
self.infinity_stk.append([]) # 末尾放一个空列表
self.infinity_stk[self.top_index].append(val) # push操作
if self.capacity > 1 : # 如果cap为1 一上来就填满了 未满set不会纳入top_index; 否则纳入top_index
self.not_full_index.add(self.top_index)
self.top_index += 1 # 继续指向下一个空栈位置
else :
first_not_full = self.not_full_index[0] # 得到第一个未空栈
self.infinity_stk[first_not_full].append(val) # push
if len(self.infinity_stk[first_not_full]) == self.capacity : # 这个栈满了
self.not_full_index.remove(first_not_full) # 将这个栈移除未满set
def pop(self) -> int: # 从右向左查看
while self.top_index > 0 and len(self.infinity_stk[self.top_index-1]) == 0 : # top_index此时未必指向右侧第一个空栈，因为popAtStack可能把中间stk弹空了
self.not_full_index.remove(self.top_index-1) # 空栈不在未满set的考虑范围内
self.top_index -= 1 # 左移top_index
if self.top_index == 0 : # 无限栈没有任何元素
return -1
ret = self.infinity_stk[self.top_index-1].pop() # 右侧第一个有元素的栈self.top_index - 1
self.not_full_index.add(self.top_index-1) # 此时self.top_index-1必为未满，可能为空
return ret
def popAtStack(self, index: int) -> int:
if index >= len(self.infinity_stk) or index < 0 : # invalid case
return -1
stk = self.infinity_stk[index]
if not stk and len(stk) == 0 : # invalid case
return -1
ret = stk.pop()
self.not_full_index.add(index) # 此时index必为一个未满栈，可能为空
return ret
# Your DinnerPlates object will be instantiated and called as such:
# obj = DinnerPlates(capacity)
# obj.push(val)
# param_2 = obj.pop()
# param_3 = obj.popAtStack(index)
