Leetcode刷题记录——面试题 03.03. 堆盘子

class StackOfPlates:

    def __init__(self, cap: int):
        self.a = []
        self.lengthlist = []
        self.limit = cap


    def push(self, val: int) -> None:
        if self.limit == 0:
            return 0
        elif self.a == []:#全空
            templist = [val]
            #self.lengthlist
            self.a.append(templist)
        elif len(self.a[-1]) < self.limit:
            self.a[-1].append(val)
        else:
            templist = [val]
            self.a.append(templist)

    def pop(self) -> int:
        
        if self.a == [] or self.limit == 0:
            return -1
        elif len(self.a[-1]) > 1:
            return self.a[-1].pop(-1)
        else:
            res = self.a[-1][0]
            self.a.pop(-1)
            return res


    def popAt(self, index: int) -> int:
        if self.limit == 0 or index >= len(self.a) or self.a[index] == []:
            return -1
        else:
            res = self.a[index].pop(-1)
            if self.a[index] == []:
                self.a.pop(index)
            return res
 


# Your StackOfPlates object will be instantiated and called as such:
# obj = StackOfPlates(cap)
# obj.push(val)
# param_2 = obj.pop()
# param_3 = obj.popAt(index)