Flatten Nested List Iterator — Stack-Based Lazy Expansion

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem 270 · Flatten Nested List Iterator

Difficulty: Medium · Pattern: Stack Design

Implement hasNext() and next() for a nested list iterator.

Solutions

# Python — stack approach
class NestedIterator:
    def __init__(self, nestedList):
        self.stack = list(reversed(nestedList))

    def next(self):
        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(reversed(top.getList()))
        return False
// Java
class NestedIterator implements Iterator<Integer> {
    Deque<NestedInteger> stack = new ArrayDeque<>();
    public NestedIterator(List<NestedInteger> nestedList) {
        for (int i=nestedList.size()-1;i>=0;i--) stack.push(nestedList.get(i));
    }
    public Integer next() { return stack.pop().getInteger(); }
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            NestedInteger top=stack.peek();
            if (top.isInteger()) return true;
            stack.pop();
            List<NestedInteger> list=top.getList();
            for (int i=list.size()-1;i>=0;i--) stack.push(list.get(i));
        }
        return false;
    }
}

Complexity

  • Time: O(1) amortized per call
  • Space: O(depth * width)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro