341. Flatten Nested List Iterator

Published on May 31, 2023

Question

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 the following 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, then your code will be judged as correct.

Example:

Input: nestedList = [1,[4,[6]]]
Output: [1,4,6]
Explanation: By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,4,6].

Algorithm

So first the hint was stack and I tried use stack, however the sequence is reversed, so I was stuck.

The intuition is quite simple. You got reverse the answer when you popping the elements, then if you want a normal order, you push the elements in a reverse order then you could get it right.

So the algorithm for this question is:

Put all the elements into the stack in reverse order;
Pop elements from stack, if it's integer, put it into the result sequence;
Other it's a list type item, we loop it and push all the elements into stack in reverse order;
Do step2-3 until stack is empty.

Well there is another way to do this where you do not need an extra list, one stack is enough.

The result list is not contracted in the construction function. You push the manipulation into the next or hasNext function like code 2.

Code

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * public 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> {
    
    Stack<NestedInteger> stack;
    List<Integer> list;
    Iterator<Integer> iter;

    public NestedIterator(List<NestedInteger> nestedList) {
        stack = new Stack<>();
        list = new ArrayList<>();
        for (int i = nestedList.size() - 1; i >= 0; i--) {
            stack.push(nestedList.get(i));
        }
        
        while (!stack.isEmpty()) {
            NestedInteger item = stack.pop();
            if (item.isInteger()) {
                list.add(item.getInteger());
            } else {
                List<NestedInteger> innerList = item.getList();
                for (int i = innerList.size() - 1; i >= 0; i--) {
                    stack.push(innerList.get(i));
                }
            }
        }
        iter = list.iterator();
    }

    @Override
    public Integer next() {
        return iter.next();
    }

    @Override
    public boolean hasNext() {
        return iter.hasNext();
    }
}

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

Code2

public class NestedIterator implements Iterator<Integer> {
    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() {
        return stack.pop().getInteger();
    }

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