Flatten Nested Iterator – 341. LeetCode


In the Flatten Nested Iterator problem, you’re given a complex data structure that contains integers or nested lists of integers (like an onion with multiple layers). Your task is to design an iterator that can “flatten” this structure so you can iterate through the integers as if they were in a single flat list. This problem is about handling nested data structures and is commonly seen in scenarios involving tree-like or recursive data structures.

Brute Force Approach

The brute force approach involves fully flattening the nested structure into a single list during initialization. Then, the iterator simply traverses this pre-flattened list.

import java.util.*;

public class NestedIterator implements Iterator<Integer> {
    private List<Integer> flatList;
    private int index;

    public NestedIterator(List<NestedInteger> nestedList) {
        flatList = new ArrayList<>();
        index = 0;
        flatten(nestedList);
    }

    private void flatten(List<NestedInteger> nestedList) {
        for (NestedInteger ni : nestedList) {
            if (ni.isInteger()) {
                flatList.add(ni.getInteger());
            } else {
                flatten(ni.getList());
            }
        }
    }

    @Override
    public Integer next() {
        return flatList.get(index++);
    }

    @Override
    public boolean hasNext() {
        return index < flatList.size();
    }
}

1. Initialization:

• Flatten the nested list into a single list (flatList) using recursion.

2. Iteration:

• next() retrieves the next integer from the flattened list.

• hasNext() checks if there are remaining integers.

Space Complexity: O(n), where n is the total number of integers because the entire structure is pre-flattened.

Time Complexity: The flattening process happens during initialization, taking O(n).

Optimized Approach: On-Demand Flattening

The optimized approach avoids fully flattening the nested list upfront. Instead, it uses a stack to traverse the structure on demand. The stack ensures that only the necessary parts of the list are processed when the iterator is used.

Algorithm and Data Structure

Algorithm: Use a stack to simulate recursive traversal of the nested structure.

Data Structure: A stack to hold the current position in the nested structure.

import java.util.*;

public class NestedIterator implements Iterator<Integer> {
    private Stack<NestedInteger> stack;

    public NestedIterator(List<NestedInteger> nestedList) {
        stack = new Stack<>();
        // Push the entire list into the stack in reverse order
        for (int i = nestedList.size() - 1; i >= 0; i--) {
            stack.push(nestedList.get(i));
        }
    }

    @Override
    public Integer next() {
        // Ensure the top is an integer
        hasNext();
        return stack.pop().getInteger();
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            NestedInteger top = stack.peek();
            if (top.isInteger()) {
                return true; // Top is an integer
            }
            // Top is a list; unpack it
            stack.pop();
            for (int i = top.getList().size() - 1; i >= 0; i--) {
                stack.push(top.getList().get(i));
            }
        }
        return false; // Stack is empty
    }
}

Step 1: Initialize the stack

public NestedIterator(List<NestedInteger> nestedList) {
    stack = new Stack<>();
    for (int i = nestedList.size() - 1; i >= 0; i--) {
        stack.push(nestedList.get(i));
    }
}


• Initializes the stack by pushing all elements of the input nestedList onto the stack in reverse order.
• The stack represents the current position in the traversal.
• Aligns with the iterator pattern by deferring computations to the point of access (next() and hasNext()).

Step 3: next() method

@Override
public Integer next() {
    hasNext(); // Ensure the top is an integer
    return stack.pop().getInteger();
}

• The hasNext() call ensures that next() only operates on integers, unpacking any nested lists if necessary.

• Guarantees the iterator behaves correctly, returning integers in sequence.

Step 4: hasNext() method

@Override
public boolean hasNext() {
    while (!stack.isEmpty()) {
        NestedInteger top = stack.peek();
        if (top.isInteger()) {
            return true; // Top is an integer
        }
        // Top is a list; unpack it
        stack.pop();
        for (int i = top.getList().size() - 1; i >= 0; i--) {
            stack.push(top.getList().get(i));
        }
    }
    return false; // Stack is empty
}

• Ensures the iterator flattens only the necessary portions of the nested list structure.

• Guarantees next() works correctly by preparing the stack.

• Ensures the iterator flattens only the necessary portions of the nested list structure.

• Guarantees next() works correctly by preparing the stack.

Why This Problem Is Important


This problem is an excellent test of a candidate’s ability to handle nested data structures, design lazy evaluation mechanisms, and use stack-based traversal effectively. It demonstrates a clear understanding of algorithmic efficiency, space-time trade-offs, and real-world applicability, especially in scenarios involving tree-like data structures.

Leave a Reply

Your email address will not be published. Required fields are marked *