Taro Logo

Flatten Nested List Iterator

Medium
2 views
2 months ago

You are given a nested list of integers nestedList. Each element is either an integer or a list whose elements may also be integers or other lists. Implement an iterator to flatten it. Implement the NestedIterator class: NestedIterator(List<NestedInteger> nestedList) initializes the iterator with the nested list nestedList. int next() returns the next integer in the nested list. boolean hasNext() returns true if there are still some integers in the nested list and false otherwise. Your code will be tested with pseudocode: initialize iterator with nestedList, res = [], while iterator.hasNext(), append iterator.next() to the end of res, return res. If res matches the expected flattened list, your code is correct. Example 1: Input: nestedList = [[1,1],2,[1,1]], Output: [1,1,2,1,1]. Example 2: Input: nestedList = [1,[4,[6]]], Output: [1,4,6]. Constraints: 1 <= nestedList.length <= 500. The values of the integers in the nested list is in the range [-10^6, 10^6].

Sample Answer
import java.util.Iterator;
import java.util.List;
import java.util.Stack;

// This is the interface that allows for creating nested lists.
// You should not implement it, or speculate about its implementation
interface NestedInteger {

    // @return true if this NestedInteger holds a single integer, rather than a nested list.
    public boolean isInteger();

    // @return the single integer that this NestedInteger holds, if it holds a single integer
    // Return null if this NestedInteger holds a nested list
    public Integer getInteger();

    // @return the nested list that this NestedInteger holds, if it holds a nested list
    // Return empty list if this NestedInteger holds a single integer
    public List<NestedInteger> getList();
}

public class NestedIterator implements Iterator<Integer> {

    private Stack<NestedInteger> stack;

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

    @Override
    public Integer next() {
        if (!hasNext()) {
            return null; // Or throw an exception
        }
        return stack.pop().getInteger();
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            NestedInteger top = stack.peek();
            if (top.isInteger()) {
                return true;
            } else {
                stack.pop();
                List<NestedInteger> list = top.getList();
                for (int i = list.size() - 1; i >= 0; i--) {
                    stack.push(list.get(i));
                }
            }
        }
        return false;
    }
}

/**
 * Your NestedIterator object will be instantiated and called as such:
 * NestedIterator i = new NestedIterator(nestedList);
 * while (i.hasNext()) v[f()] = i.next();
 */

Naive Solution

A naive approach would be to flatten the entire nested list into a single list when the iterator is initialized. This would make the next() and hasNext() methods very simple, but it would require extra space to store the flattened list and would not be efficient if only a small portion of the list is iterated.

Optimal Solution

The optimal solution uses a stack to keep track of the nested lists. The hasNext() method checks if the stack is empty. If it's not empty, it peeks at the top element. If the top element is an integer, then hasNext() returns true. If the top element is a list, then the list is unrolled and its elements are pushed onto the stack in reverse order. This ensures that the elements are processed in the correct order. The next() method simply pops the top element from the stack and returns its integer value.

Example

Consider the nested list [[1,1],2,[1,1]].

  1. The NestedIterator is initialized with this list.
  2. The stack is initialized with the elements in reverse order: [1,1], 2, [1,1].
  3. hasNext() is called. The stack is not empty. The top element is [1,1], which is a list.
  4. The list [1,1] is unrolled and its elements are pushed onto the stack in reverse order: 1, 1, 2, [1,1].
  5. hasNext() is called. The stack is not empty. The top element is 1, which is an integer. hasNext() returns true.
  6. next() is called. The top element 1 is popped from the stack and returned.
  7. The stack is now 1, 2, [1,1].
  8. This process continues until all elements have been iterated.

Big(O) Run-time Analysis

  • hasNext(): O(1) on average. While it might take O(N) in the worst case to flatten nested lists at the top of the stack, the amortized time complexity is O(1) because each nested list and integer is visited at most twice (once when pushed onto the stack and once when popped).
  • next(): O(1) on average. Similar to hasNext(), the amortized time complexity is O(1).

Big(O) Space Usage Analysis

  • O(D), where D is the maximum depth of the nested list. This is because, in the worst case, the stack will contain all the nested lists at the maximum depth of the nested structure.

Edge Cases

  • Empty List: If the input list is empty, the hasNext() method should return false immediately.
  • Nested Empty Lists: If the input list contains nested empty lists, the iterator should skip them without returning any elements.
  • Large Depth: The solution should handle arbitrarily deep nested lists without causing a stack overflow. The stack-based approach is suitable for this.
  • Null elements: The code assumes that NestedInteger only holds an Integer or List. If a NestedInteger can hold a null value, we should add null checks in the hasNext() and next() methods to handle them correctly, possibly by skipping them or throwing an exception, depending on the specific requirements.