程序员面试金典 - 面试题 03.03. 堆盘子

题目难度: 中等

原题链接

今天继续更新程序员面试金典系列, 大家在公众号 算法精选 里回复 面试金典 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

堆盘子。设想有一堆盘子，堆太高可能会倒下来。因此，在现实生活中，盘子堆到一定高度时，我们就会另外堆一堆盘子。请实现数据结构 SetOfStacks，模拟这种行为。SetOfStacks 应该由多个栈组成，并且在前一个栈填满时新建一个栈。此外，SetOfStacks.push()和 SetOfStacks.pop()应该与普通栈的操作方法相同（也就是说，pop()返回的值，应该跟只有一个栈时的情况一样）。 进阶：实现一个 popAt(int index)方法，根据指定的子栈，执行 pop 操作。

当某个栈为空时，应当删除该栈。当栈中没有元素或不存在该栈时，pop，popAt 应返回 -1.

示例 1:

• 输入：
  ["StackOfPlates", "push", "push", "popAt", "pop", "pop"]
  [[1], [1], [2], [1], [], []]
• 输出：
  [null, null, null, 2, 1, -1]

示例 2:

• 输入：
  ["StackOfPlates", "push", "push", "push", "popAt", "popAt", "popAt"]
  [[2], [1], [2], [3], [0], [0], [0]]
• 输出：
  [null, null, null, null, 2, 1, 3]

题目思考

1. 需要存储哪些信息, 如何操作?

解决方案

思路

• 根据题目描述, 我们可以直接构造存储多个栈的数组, 并需要额外保存栈的上限
• push 操作: 如果数组为空或最后一个栈满了, 则新增一个栈并存当前元素, 否则直接追加到最后一个栈上
• popAt 操作: 先判断下标是否有效, 无效返回-1, 有效时则 pop 对应栈的栈顶, 如果 pop 后栈为空则从数组中移除这个栈
• pop 操作: 可以直接利用 popAt(len-1), 简化代码

复杂度

• 时间复杂度
  • push 和 pop 显然是 O(1)
  • popAt 大部分情况也是 O(1), 只有栈变空后才是 O(N)
• 空间复杂度 O(N)
  • 使用了若干个栈存 N 个元素

代码

class StackOfPlates:
    # 简单模拟
    def __init__(self, cap: int):
        self.cap = cap
        self.stacks = []

    def push(self, val: int) -> None:
        # 注意如果cap为0则不能push
        if self.cap == 0:
            return
        if not self.stacks or len(self.stacks[-1]) == self.cap:
            # 没有栈或者最后一个栈满了时, 需要新增加一个栈
            self.stacks.append([val])
        else:
            # 最后一个栈不满, 则可以直接追加
            self.stacks[-1].append(val)

    def pop(self) -> int:
        # 可以直接调用popAt(len-1)
        return self.popAt(len(self.stacks) - 1)

    def popAt(self, index: int) -> int:
        if index < 0 or index >= len(self.stacks):
            # 不合法的index, 返回-1
            # 注意这里无需判断对应index的栈是否为空, 因为某个栈pop后为空的话一定会被删除, 不可能存在空的栈
            return -1
        res = self.stacks[index].pop()
        if not self.stacks[index]:
            # 如果pop后栈为空的话, 需要删除这个栈
            del self.stacks[index]
        return res

大家可以在下面这些地方找到我~😊

我的知乎专栏

我的头条号

我的 CSDN

我的 Leetcode

我的牛客网博客

我的公众号: 算法精选, 欢迎大家扫码关注~😊