LeetCode: Flatten Nested List Iterator

LeetCode: Flatten Nested List Iterator

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:

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

Example 2:

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

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
26
27
28
29
30
31
32
public class NestedIterator implements Iterator<Integer> {

  private Stack<NestedInteger> stack;

  public NestedIterator(List<NestedInteger> nestedList) {
      stack = new Stack<>();
      if (nestedList == null || nestedList.isEmpty()) return;
      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 ni = stack.peek();
          if (ni.isInteger()) return true;
          stack.pop();
          List<NestedInteger> list = ni.getList();
          for (int i = list.size() - 1; i >= 0; i--) {
              stack.push(list.get(i));
          }
      }
      return false;
  }
}