[LintCode/LeetCode] Flatten Nested List Iterator

444 查看

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()判断下一个NestedIntegerInteger还是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();
    }
}