341. 扁平化嵌套列表迭代器
题目描述
输入一个嵌套整数列表 nestedList。列表里每一项要么是整数,要么是另一个列表,而这个内层列表的每一项同样可能是整数或列表,层数不限。
请实现迭代器类 NestedIterator,按从左到右的顺序依次给出所有整数:
NestedIterator(nestedList):用给定的嵌套列表创建迭代器。next():返回下一个整数。hasNext():还有没取出的整数时返回 true,否则返回 false。
评测方式:创建迭代器后,只要 hasNext() 为 true 就调用 next(),把结果依次放进列表 res。如果 res 等于把嵌套列表展平后的结果,就算正确。示例的输出就是这个 res。
输入:nestedList = [[1,1],2,[1,1]] 输出:[1,1,2,1,1]
解释:反复调用 next(),依次得到 1,1,2,1,1。
输入: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 Falseclass NestedIterator implements Iterator<Integer> {
private Deque<NestedInteger> st = new ArrayDeque<>();
public NestedIterator(List<NestedInteger> nestedList) {
for (int i = nestedList.size() - 1; i >= 0; i--) {
st.push(nestedList.get(i));
}
}
@Override
public Integer next() {
hasNext();
return st.pop().getInteger();
}
@Override
public boolean hasNext() {
while (!st.isEmpty()) {
NestedInteger top = st.peek();
if (top.isInteger()) {
return true;
}
st.pop();
List<NestedInteger> ch = top.getList();
for (int i = ch.size() - 1; i >= 0; i--) {
st.push(ch.get(i));
}
}
return false;
}
}class NestedIterator {
vector<NestedInteger> st;
public:
NestedIterator(vector<NestedInteger> &nestedList) {
st.assign(nestedList.rbegin(), nestedList.rend());
}
int next() {
hasNext();
int v = st.back().getInteger();
st.pop_back();
return v;
}
bool hasNext() {
while (!st.empty()) {
NestedInteger top = st.back();
if (top.isInteger()) {
return true;
}
st.pop_back();
const vector<NestedInteger> &ch = top.getList();
for (int i = (int)ch.size() - 1; i >= 0; i--) {
st.push_back(ch[i]);
}
}
return false;
}
};type NestedIterator struct {
st []*NestedInteger
}
func Constructor(nestedList []*NestedInteger) *NestedIterator {
st := []*NestedInteger{}
for i := len(nestedList) - 1; i >= 0; i-- {
st = append(st, nestedList[i])
}
return &NestedIterator{st: st}
}
func (this *NestedIterator) Next() int {
this.HasNext()
top := this.st[len(this.st)-1]
this.st = this.st[:len(this.st)-1]
return top.GetInteger()
}
func (this *NestedIterator) HasNext() bool {
for len(this.st) > 0 {
top := this.st[len(this.st)-1]
if top.IsInteger() {
return true
}
this.st = this.st[:len(this.st)-1]
ch := top.GetList()
for i := len(ch) - 1; i >= 0; i-- {
this.st = append(this.st, ch[i])
}
}
return false
}var NestedIterator = function(nestedList) {
this.st = nestedList.slice().reverse();
};
NestedIterator.prototype.next = function() {
this.hasNext();
return this.st.pop().getInteger();
};
NestedIterator.prototype.hasNext = function() {
while (this.st.length > 0) {
const top = this.st[this.st.length - 1];
if (top.isInteger()) {
return true;
}
this.st.pop();
const ch = top.getList();
for (let i = ch.length - 1; i >= 0; i--) {
this.st.push(ch[i]);
}
}
return false;
};class NestedIterator {
st: NestedInteger[];
constructor(nestedList: NestedInteger[]) {
this.st = nestedList.slice().reverse();
}
next(): number {
this.hasNext();
return this.st.pop()!.getInteger()!;
}
hasNext(): boolean {
while (this.st.length > 0) {
const top = this.st[this.st.length - 1];
if (top.isInteger()) {
return true;
}
this.st.pop();
const ch = top.getList();
for (let i = ch.length - 1; i >= 0; i--) {
this.st.push(ch[i]);
}
}
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)class NestedIterator implements Iterator<Integer> {
private List<Integer> vals = new ArrayList<>();
private int pos = 0;
public NestedIterator(List<NestedInteger> nestedList) {
dfs(nestedList);
}
private void dfs(List<NestedInteger> lst) {
for (NestedInteger ni : lst) {
if (ni.isInteger()) vals.add(ni.getInteger());
else dfs(ni.getList());
}
}
@Override
public Integer next() {
return vals.get(pos++);
}
@Override
public boolean hasNext() {
return pos < vals.size();
}
}class NestedIterator {
vector<int> vals;
int pos = 0;
void dfs(const vector<NestedInteger> &lst) {
for (const NestedInteger &ni : lst) {
if (ni.isInteger()) vals.push_back(ni.getInteger());
else dfs(ni.getList());
}
}
public:
NestedIterator(vector<NestedInteger> &nestedList) {
dfs(nestedList);
}
int next() {
return vals[pos++];
}
bool hasNext() {
return pos < (int)vals.size();
}
};type NestedIterator struct {
vals []int
pos int
}
func Constructor(nestedList []*NestedInteger) *NestedIterator {
vals := []int{}
var dfs func(lst []*NestedInteger)
dfs = func(lst []*NestedInteger) {
for _, ni := range lst {
if ni.IsInteger() {
vals = append(vals, ni.GetInteger())
} else {
dfs(ni.GetList())
}
}
}
dfs(nestedList)
return &NestedIterator{vals: vals}
}
func (this *NestedIterator) Next() int {
v := this.vals[this.pos]
this.pos++
return v
}
func (this *NestedIterator) HasNext() bool {
return this.pos < len(this.vals)
}var NestedIterator = function(nestedList) {
this.vals = [];
this.pos = 0;
const dfs = (lst) => {
for (const ni of lst) {
if (ni.isInteger()) this.vals.push(ni.getInteger());
else dfs(ni.getList());
}
};
dfs(nestedList);
};
NestedIterator.prototype.next = function() {
return this.vals[this.pos++];
};
NestedIterator.prototype.hasNext = function() {
return this.pos < this.vals.length;
};class NestedIterator {
vals: number[] = [];
pos = 0;
constructor(nestedList: NestedInteger[]) {
this.dfs(nestedList);
}
dfs(lst: NestedInteger[]): void {
for (const ni of lst) {
if (ni.isInteger()) this.vals.push(ni.getInteger()!);
else this.dfs(ni.getList());
}
}
next(): number {
return this.vals[this.pos++];
}
hasNext(): boolean {
return this.pos < this.vals.length;
}
}