在问 · 算法题解

341. 扁平化嵌套列表迭代器

中等设计栈(惰性展开)力扣原题 ↗

题目描述

输入一个嵌套整数列表 nestedList。列表里每一项要么是整数,要么是另一个列表,而这个内层列表的每一项同样可能是整数或列表,层数不限。

请实现迭代器类 NestedIterator,按从左到右的顺序依次给出所有整数:

NestedIterator(nestedList):用给定的嵌套列表创建迭代器。next():返回下一个整数。hasNext():还有没取出的整数时返回 true,否则返回 false。

评测方式:创建迭代器后,只要 hasNext() 为 true 就调用 next(),把结果依次放进列表 res。如果 res 等于把嵌套列表展平后的结果,就算正确。示例的输出就是这个 res。

示例 1
输入:nestedList = [[1,1],2,[1,1]]
输出:[1,1,2,1,1]

解释:反复调用 next(),依次得到 1,1,2,1,1。

示例 2
输入:nestedList = [1,[4,[6]]]
输出:[1,4,6]

解释:逐层展开后依次得到 1,4,6。

数据范围

  • 1 <= nestedList.length <= 500
  • 嵌套列表中所有整数都在 [-10^6, 10^6] 范围内

解题思路:栈(惰性展开)

用栈模拟递归,并且按需展开:构造时把顶层元素逆序压栈。hasNext 循环查看栈顶,栈顶是整数就返回 True;是列表就弹出,再把它的子元素逆序压栈,然后继续看新栈顶;栈空时返回 False。next 先调用 hasNext,保证栈顶是整数,再把它弹出返回。

为什么正确

栈中元素从栈顶往下,按原列表从左到右的顺序排列;hasNext 返回 True 之后,栈顶一定是下一个要产出的整数。

复杂度分析

设整数和列表总数为 N,每个元素最多入栈、出栈一次,所有调用合计 O(N),单次均摊 O(1);栈的空间最多 O(N)。

易错点

只判断“栈非空”不够:栈里可能只剩空列表,比如 [[],[[]]],这时必须真正拆到底才能确定没有整数。所以展开逻辑要放在 hasNext 里,不能放在 next 里。另外子元素要逆序压栈,否则输出顺序会反过来。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。

class NestedIterator:
    def __init__(self, nestedList):
        self.stack = nestedList[::-1]

    def next(self):
        self.hasNext()
        return self.stack.pop().getInteger()

    def hasNext(self):
        while self.stack:
            top = self.stack[-1]
            if top.isInteger():
                return True
            self.stack.pop()
            self.stack.extend(top.getList()[::-1])
        return False

载入编辑器,自己跑一遍 →

其他解法

预先展平把全部工作放在构造阶段;栈的写法按需展开,只在调用 hasNext 时处理到下一个整数,更符合迭代器的惰性语义。两者总时间都是 O(n)。

解法 2:深度优先搜索(预先展平) 时间 初始化 O(n),next/hasNext O(1),空间 O(n)

构造时对嵌套列表做一次深度优先搜索。遇到整数就追加到数组,遇到列表就递归展开。之后用一个下标遍历这个数组:next 返回当前元素并后移,hasNext 判断下标是否越界。

实现最直观,next 和 hasNext 都是严格 O(1)。缺点是构造时要一次性展开全部元素,即使只读前几个也要付出完整的时间和空间,且递归深度等于嵌套深度。

class NestedIterator:
    def __init__(self, nestedList):
        self.vals = []
        self.pos = 0
        self.dfs(nestedList)

    def dfs(self, lst):
        for ni in lst:
            if ni.isInteger():
                self.vals.append(ni.getInteger())
            else:
                self.dfs(ni.getList())

    def next(self):
        v = self.vals[self.pos]
        self.pos += 1
        return v

    def hasNext(self):
        return self.pos < len(self.vals)

设计 专题的其他题目

查看 设计 全部题目 →