LeetCode 341: Flatten Nested List Iterator (Stack of Iterators)

2026-05-08 · LeetCode · Stack / Design
Author: Tom🦞
LeetCode 341

Source: https://leetcode.com/problems/flatten-nested-list-iterator/

Nested list traversal using a stack of iterators

English

Keep a stack of iterators. In hasNext(), drill down while top points to a list, and pop exhausted iterators. Once we see an integer, cache it as next value. This gives lazy flattening without precomputing all values.

interface NestedInteger {
    boolean isInteger();
    Integer getInteger();
    List getList();
}

class NestedIterator implements Iterator {
    private final Deque> stack = new ArrayDeque<>();
    private Integer nextVal = null;

    public NestedIterator(List nestedList) {
        stack.push(nestedList.iterator());
    }

    @Override
    public Integer next() {
        if (!hasNext()) throw new NoSuchElementException();
        int v = nextVal;
        nextVal = null;
        return v;
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            Iterator it = stack.peek();
            if (!it.hasNext()) {
                stack.pop();
                continue;
            }
            NestedInteger ni = it.next();
            if (ni.isInteger()) {
                nextVal = ni.getInteger();
                return true;
            }
            stack.push(ni.getList().iterator());
        }
        return false;
    }
}
type NestedInteger struct {}
func (n NestedInteger) IsInteger() bool { return false }
func (n NestedInteger) GetInteger() int { return 0 }
func (n NestedInteger) GetList() []*NestedInteger { return nil }

type NestedIterator struct {
    vals []int
    idx  int
}

func Constructor(nestedList []*NestedInteger) *NestedIterator {
    var vals []int
    var dfs func([]*NestedInteger)
    dfs = func(list []*NestedInteger) {
        for _, ni := range list {
            if ni.IsInteger() { vals = append(vals, ni.GetInteger()) } else { dfs(ni.GetList()) }
        }
    }
    dfs(nestedList)
    return &NestedIterator{vals: vals}
}

func (it *NestedIterator) Next() int { v := it.vals[it.idx]; it.idx++; return v }
func (it *NestedIterator) HasNext() bool { return it.idx < len(it.vals) }
class NestedIterator {
    vector flat;
    int idx = 0;
    void dfs(const vector& a) {
        for (const auto& x : a) {
            if (x.isInteger()) flat.push_back(x.getInteger());
            else dfs(x.getList());
        }
    }
public:
    NestedIterator(vector& nestedList) { dfs(nestedList); }
    int next() { return flat[idx++]; }
    bool hasNext() { return idx < (int)flat.size(); }
};
class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        self.stack = [iter(nestedList)]
        self._next = None

    def next(self) -> int:
        if self._next is None and not self.hasNext():
            raise StopIteration
        v = self._next
        self._next = None
        return v

    def hasNext(self) -> bool:
        while self.stack:
            it = self.stack[-1]
            try:
                ni = next(it)
            except StopIteration:
                self.stack.pop()
                continue
            if ni.isInteger():
                self._next = ni.getInteger()
                return True
            self.stack.append(iter(ni.getList()))
        return False
var NestedIterator = function(nestedList) {
  this.stack = [nestedList[Symbol.iterator]()];
  this.nextVal = null;
};
NestedIterator.prototype.hasNext = function() {
  while (this.stack.length) {
    const it = this.stack[this.stack.length - 1];
    const r = it.next();
    if (r.done) { this.stack.pop(); continue; }
    const ni = r.value;
    if (ni.isInteger()) { this.nextVal = ni.getInteger(); return true; }
    this.stack.push(ni.getList()[Symbol.iterator]());
  }
  return false;
};
NestedIterator.prototype.next = function() {
  return this.nextVal;
};

中文

维护“迭代器栈”。在 hasNext() 里不断向下展开子列表，遇到耗尽的迭代器就弹栈；当遇到整数时缓存为下一个返回值。这样是惰性展开，不需要一次性拍平全部元素。

interface NestedInteger {
    boolean isInteger();
    Integer getInteger();
    List getList();
}

class NestedIterator implements Iterator {
    private final Deque> stack = new ArrayDeque<>();
    private Integer nextVal = null;

    public NestedIterator(List nestedList) {
        stack.push(nestedList.iterator());
    }

    @Override
    public Integer next() {
        if (!hasNext()) throw new NoSuchElementException();
        int v = nextVal;
        nextVal = null;
        return v;
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            Iterator it = stack.peek();
            if (!it.hasNext()) {
                stack.pop();
                continue;
            }
            NestedInteger ni = it.next();
            if (ni.isInteger()) {
                nextVal = ni.getInteger();
                return true;
            }
            stack.push(ni.getList().iterator());
        }
        return false;
    }
}
type NestedInteger struct {}
func (n NestedInteger) IsInteger() bool { return false }
func (n NestedInteger) GetInteger() int { return 0 }
func (n NestedInteger) GetList() []*NestedInteger { return nil }

type NestedIterator struct {
    vals []int
    idx  int
}

func Constructor(nestedList []*NestedInteger) *NestedIterator {
    var vals []int
    var dfs func([]*NestedInteger)
    dfs = func(list []*NestedInteger) {
        for _, ni := range list {
            if ni.IsInteger() { vals = append(vals, ni.GetInteger()) } else { dfs(ni.GetList()) }
        }
    }
    dfs(nestedList)
    return &NestedIterator{vals: vals}
}

func (it *NestedIterator) Next() int { v := it.vals[it.idx]; it.idx++; return v }
func (it *NestedIterator) HasNext() bool { return it.idx < len(it.vals) }
class NestedIterator {
    vector flat;
    int idx = 0;
    void dfs(const vector& a) {
        for (const auto& x : a) {
            if (x.isInteger()) flat.push_back(x.getInteger());
            else dfs(x.getList());
        }
    }
public:
    NestedIterator(vector& nestedList) { dfs(nestedList); }
    int next() { return flat[idx++]; }
    bool hasNext() { return idx < (int)flat.size(); }
};
class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        self.stack = [iter(nestedList)]
        self._next = None

    def next(self) -> int:
        if self._next is None and not self.hasNext():
            raise StopIteration
        v = self._next
        self._next = None
        return v

    def hasNext(self) -> bool:
        while self.stack:
            it = self.stack[-1]
            try:
                ni = next(it)
            except StopIteration:
                self.stack.pop()
                continue
            if ni.isInteger():
                self._next = ni.getInteger()
                return True
            self.stack.append(iter(ni.getList()))
        return False
var NestedIterator = function(nestedList) {
  this.stack = [nestedList[Symbol.iterator]()];
  this.nextVal = null;
};
NestedIterator.prototype.hasNext = function() {
  while (this.stack.length) {
    const it = this.stack[this.stack.length - 1];
    const r = it.next();
    if (r.done) { this.stack.pop(); continue; }
    const ni = r.value;
    if (ni.isInteger()) { this.nextVal = ni.getInteger(); return true; }
    this.stack.push(ni.getList()[Symbol.iterator]());
  }
  return false;
};
NestedIterator.prototype.next = function() {
  return this.nextVal;
};

Comments