341. Flatten Nested List Iterator

Description

The NestedIterator problem involves flattening a nested list of integers. A nested list can contain integers or other lists, and the goal is to design an iterator that returns all integers from the nested structure in a sequential manner. The iterator should support two operations: next() to get the next integer, and hasNext() to check if there are more integers to iterate over.

Intuition

The nested list structure can be represented as a tree where each integer is a leaf, and each list is a node that might contain other nodes or leaves. The key insight is to simulate a depth-first traversal of this tree, yielding integers whenever encountered. A stack or recursion can help in traversing the nested structure, keeping track of where we are in the nested lists.

Approach

1. Initialization: Start by pushing the iterator of the initial nested list onto a stack (for iterative solutions). Each item can either be an integer or another list.
2. hasNext(): Check if the current element in the list is an integer or a nested list. If it’s an integer, return true. If it’s a nested list, push the iterator for the nested list onto the stack and continue checking. Pop the stack when a sublist is exhausted.
3. next(): Assume hasNext() has been called, ensuring the iterator is at a valid integer. Retrieve and return the current integer and advance the iterator to the next position.
4. Recursive Approach (Python): The recursion keeps track of all nested lists by generating integers one by one, leveraging Python's yield to flatten the structure.

Complexity

Time Complexity:

The hasNext() and next() operations together visit each integer exactly once, thus the time complexity is proportional to the number of integers in the nested list, i.e., O(n), where n is the total number of integers.

Space Complexity:

The space complexity depends on the depth of the nested lists. In the worst case, if all the lists are deeply nested, the space complexity would be O(d), where d is the maximum depth of the nested list. In case of iterative solutions, it’s proportional to the depth of the list structure.

Code

C++

class NestedIterator {
    stack::iterator> begins, ends;
    
public:
    NestedIterator(vector &nestedList) {
        begins.push(nestedList.begin());
        ends.push(nestedList.end());
    }

    int next() {
        hasNext();  // Ensure the iterator is at a valid integer.
        return (begins.top()++)->getInteger();
    }

    bool hasNext() {
        while (!begins.empty()) {
            if (begins.top() == ends.top()) {  // Current list is fully traversed
                begins.pop();
                ends.pop();
            } else if (begins.top()->isInteger()) {
                return true;
            } else {  // The current element is a list
                auto &nestedList = begins.top()->getList();
                begins.top()++;
                begins.push(nestedList.begin());
                ends.push(nestedList.end());
            }
        }
        return false;
    }
};

Python

class NestedIterator:
    def __init__(self, nestedList: [NestedInteger]):
        self.generator = self._flatten(nestedList)
        self.nextVal = None

    def _flatten(self, nestedList):
        for item in nestedList:
            if item.isInteger():
                yield item.getInteger()
            else:
                yield from self._flatten(item.getList())

    def next(self) -> int:
        if not self.hasNext():
            raise Exception("No more elements")
        next_val = self.nextVal
        self.nextVal = None
        return next_val

    def hasNext(self) -> bool:
        if self.nextVal is not None:
            return True
        try:
            self.nextVal = next(self.generator)
            return True
        except StopIteration:
            return False

Java

public class NestedIterator implements Iterator {

    private Stack> stack;

    public NestedIterator(List nestedList) {
        stack = new Stack<>();
        stack.push(nestedList.listIterator());
    }

    @Override
    public Integer next() {
        hasNext();  // Ensure there's an integer ready to return.
        return stack.peek().next().getInteger();
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            if (!stack.peek().hasNext()) {
                stack.pop();  // Finished processing the current list.
            } else {
                NestedInteger x = stack.peek().next();
                if (x.isInteger()) {
                    stack.peek().previous();  // Step back to handle it in `next()`
                    return true;
                }
                stack.push(x.getList().listIterator());
            }
        }
        return false;
    }
}

JavaScript

var NestedIterator = function (nestedList)
 {
    this.nestedList = nestedList;
    this.iterator = flattenIteratorGen(this.nestedList);
    this.currentValue = null;
};

/**
 * @this NestedIterator
 * @returns {boolean}
 */
NestedIterator.prototype.hasNext = function () {
    if (this.currentValue !== null) {
        return true;
    }
    const next = this.iterator.next();
    if (next.done) {
        return false;
    }
    this.currentValue = next.value;
    return true;
};

/**
 * @this NestedIterator
 * @returns {number}
 */
NestedIterator.prototype.next = function () {
    if (!this.hasNext()) {
        throw new Error("No more elements to iterate.");
    }

    const value = this.currentValue;
    this.currentValue = null;
    return value;
};

const flattenIteratorGen = function* (nestedList) {
    for (var item of nestedList) {
        if (item.isInteger()) {
            yield item.getInteger();
        } else {
            yield* flattenIteratorGen(item.getList());
        }
    }
};