341. Flatten Nested List Iterator

Problem¶

Implement a lazy iterator over nested integers that returns every integer in left-to-right order through next() and reports availability through hasNext().

Input:  [1, [4, [6]], []]
Calls:  while hasNext(): next()
Output: [1, 4, 6]

#class NestedInteger:
#    def isInteger(self) -> bool:
#        """
#        @return True if this NestedInteger holds a single integer, rather than a nested list.
#        """
#
#    def getInteger(self) -> int:
#        """
#        @return the single integer that this NestedInteger holds, if it holds a single integer
#        Return None if this NestedInteger holds a nested list
#        """
#
#    def getList(self) -> [NestedInteger]:
#        """
#        @return the nested list that this NestedInteger holds, if it holds a nested list
#        Return None if this NestedInteger holds a single integer
#        """

class NestedIterator:
    def __init__(self, nestedList: ["NestedInteger"]):
        pass

    def next(self) -> int:
        pass

Key trick¶

Keep a stack of list iterators; hasNext() descends through nested lists until it finds and caches an integer, making repeated calls to hasNext() safe.

Trap¶

Treating a nonempty stack as proof that an integer remains.
Leaving exhausted or out-of-range list positions on the stack.
Failing on empty nested lists.
Advancing again when hasNext() is called repeatedly before next().

Why is it interesting?¶

It tests lazy depth-first traversal, explicit recursion state, and the iterator contract without flattening the entire input first.

Python solution¶

class NestedIterator:
    def __init__(self, nestedList: list["NestedInteger"]):
        # Each iterator represents one active nesting level.
        self.stack = [iter(nestedList)]
        self.nxt = 0
        self.ready = False

    def next(self) -> int:
        if not self.hasNext():
            raise StopIteration

        self.ready = False
        return self.nxt

    def hasNext(self) -> bool:
        # Preserve a cached integer when hasNext() is called repeatedly.
        if self.ready:
            return True

        while self.stack:
            try:
                item = next(self.stack[-1])
            except StopIteration:
                self.stack.pop()
                continue

            if item.isInteger():
                self.nxt = item.getInteger()
                self.ready = True
                return True

            self.stack.append(iter(item.getList()))

        return False

Time complexity: \(O(N)\) over the full traversal, where \(N\) is the number of nested elements.

Space complexity: \(O(D)\), where \(D\) is the maximum nesting depth.

Comment on my solution¶

The stack-of-cursors idea is appropriate, but hasNext() only checks whether the stack is nonempty rather than whether a valid integer remains.

For example, [[]] fails because descending into the empty list attempts to read position 0; the parent cursor (nested_list, 1) is also out of range. Cursors after nested lists are pushed without first checking whether another parent element exists. Normalizing the traversal state inside hasNext() fixes these cases and supports repeated hasNext() calls.

# Ideas
# - [[1,1],2,[1,1]] -> [1,1,2,1,1]
# - [1,[4,[6]]] -> [1,4,6]
# - keep a pointer to next position
#   - I think we need a stack for this
# - (bad idea) pop element after consuming it so the first element in
#   the list is always the next element to be consumed

# WRONG
class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        if len(nestedList) > 0:
            # ([NestedInteger], pos)
            self.stack = [(nestedList, 0)]
        else:
            self.stack = []

    def next(self) -> int:
        if not self.hasNext():
            raise StopIteration

        nested_list, pos = self.stack.pop()
        while not nested_list[pos].isInteger():
            self.stack.append((nested_list, pos + 1))
            nested_list = nested_list[pos].getList()
            pos = 0

        n = nested_list[pos].getInteger()
        if pos < len(nested_list) - 1:
            self.stack.append((nested_list, pos + 1))

        return n

    def hasNext(self) -> bool:
        return len(self.stack) > 0