[LeetCode] Problem 341 - Flatten Nested List Iterator

Posted on | In ,

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

No.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].

No.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].

Code

1
2
3
4
5
public interface NestedInteger {
    public boolean isInteger();
    public Integer getInteger();
    public List<NestedInteger> getList();
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
private Stack<NestedInteger> stack = new Stack<>();

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

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

@Override
public boolean hasNext() {
    while (!stack.isEmpty() && !stack.peek().isInteger()) {
        NestedInteger nestedInteger = stack.pop();
        push(nestedInteger.getList());
    }

    return !stack.isEmpty();
}

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