[LintCode/LeetCode] Flatten Nested List Iterator

Problem

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

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

Given the list [1,[4,[6]]], By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,4,6].

Note

這個迭代器要實現三個功能：
初始化結構，找下一個元素，判斷是否存在下一個元素。
首先，根據迭代器須要不斷返回下一個元素，肯定用堆棧來作。
堆棧初始化數據結構，要先從後向前向堆棧壓入nestedList中的元素。
在調用next()以前，先要用hasNext()判斷下一個NestedInteger是Integer仍是List<Integer>，並進行flatten的操做：對List<Integer>要展開並順序壓入stack；對Integer直接返回true。這樣一來，調用next()的時候，返回的就必定是Integer了。

Solution

import java.util.Iterator;
public class NestedIterator implements Iterator<Integer> {
    Stack<NestedInteger> stack = new Stack();
    public NestedIterator(List<NestedInteger> nestedList) {
        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;
    }
}

By using internalNext()

public class NestedIterator implements Iterator<Integer> {
    private NestedInteger peek = null;
    private Iterator<NestedInteger> iterator;
    private Stack<Iterator<NestedInteger>> stack = new Stack();
    public NestedIterator(List<NestedInteger> nestedList) {
        iterator = nestedList.iterator();
        peek = internalNext();
    }
    private NestedInteger internalNext() {
        if (iterator.hasNext()) {
            NestedInteger i = iterator.next();
            if (i.isInteger()) return i;
            else {
                stack.add(iterator);
                iterator = i.getList().iterator();
                return internalNext();
            }
        } else if (stack.size() > 0) {
            iterator = stack.pop();
            return internalNext();
        } else return null;
    }
    @Override
    public Integer next() {
        Integer next = peek.getInteger();
        peek = internalNext();
        return next;
    }
    @Override
    public boolean hasNext() {
        return peek != null;
    }
}

A much much better method

public class NestedIterator implements Iterator<Integer> {
    private List<Integer> flatList;
    private int index;
    public NestedIterator(List<NestedInteger> nestedList) {
        flatList = new ArrayList<>();
        flatten(nestedList);
    }
    public void flatten(List<NestedInteger> nestedList) {
        for (NestedInteger i: nestedList) {
            if (i.isInteger()) flatList.add(i.getInteger());
            else flatten(i.getList());
        }
    }
    @Override
    public Integer next() {
        return hasNext ? flatList.get(index++) : null;
    }
    @Override
    public boolean hasNext() {
        return index < flatList.size();
    }
}

update 2018-11

public class NestedIterator implements Iterator<Integer> {

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

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

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