Flatten Nested List Iterator

Stack, Design

Medium

Given a nested list of integers, implement an iterator to flatten it.

Each element is either an integer, or a list -- whose elements may also be integers or other lists.

Example 1:

Input: [[1,1],2,[1,1]]
Output: [1,1,2,1,1]

Explanation: 
By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,1,2,1,1].

Example 2:

Input: [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].

Solution

Stack

Similar to Nested List Weight Sum, we can use recursion to solve it. But since we need to access each NestedInteger at a time, we will use a stack to help.

In the constructor, we push all the nestedList into the stack from back to front, so when we pop the stack, it returns the very first element.

Second, in the hasNext() function, we peek the first element in stack currently, and if it is an Integer, we will return true and pop the element. If it is a list, we will further flatten it. This is iterative version of flatting the nested list. Again, we need to iterate from the back to front of the list.

public class NestedIterator implements Iterator<Integer> {

    private Stack<NestedInteger> stack;

    public NestedIterator(List<NestedInteger> nestedList) {
        stack = new Stack<>();
        flattenList(nestedList);
    }

    @Override
    public Integer next() {
        return hasNext() ? stack.pop().getInteger() : null;
    }

    @Override
    public boolean hasNext() {
        while (!stack.isEmpty()) {
            if (stack.peek().isInteger()) return true;
            flattenList(stack.pop().getList());
        }
        return false;
    }

    private void flattenList(List<NestedInteger> list) {
        for (int i = list.size() - 1; i >= 0; i--) {
            stack.push(list.get(i));
        }
    }
}

Pre-populate - Using Queue to Store - Recursive

public class NestedIterator implements Iterator<Integer> {

    private Queue<Integer> queue = new LinkedList();

    public NestedIterator(List<NestedInteger> nestedList) {
        helper(nestedList);
    }

    private void helper(List<NestedInteger> list){
        if (list == null) {
            return;   
        }

        for (NestedInteger in: list){
            if (in.isInteger())
                queue.offer(in.getInteger());
            else {
                helper(in.getList());
            }
        }
    }

    @Override
    public Integer next() {
        if (hasNext()) {
            return queue.poll();
        } else {
            return -1;
        }
    }

    @Override
    public boolean hasNext() {
        if (!queue.isEmpty()) {
            return true;
        } else {
            return false;
        }
    }
}

预先全部flatten - Recursive

class NestedIterator implements Iterator<Integer> {

    Iterator<Integer> iterator;

    public NestedIterator(List<NestedInteger> nestedList) {
        List<Integer> nums = new ArrayList<>();
        if(nestedList != null) {
            flatten(nestedList,nums);
        }
        iterator = nums.iterator();
    }

    private void flatten(List<NestedInteger> nestedList, List<Integer> nums) {
        for(NestedInteger n : nestedList) {
            if(n.isInteger()) {
                nums.add(n.getInteger());
            } else {
                flatten(n.getList(),nums);
            }
        }
    }

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

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

Last updated