Reference: LeetCode
Difficulty: Medium

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.

Note:

Example: 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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].

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

Input: [[]]
Output: []

Input: [[[]]]
Output: []

Analysis

Use Stack

Reference: Simple Java solution using a stack with explanation

Use a stack to store all NestedInteger. At the beginning, we push into all NestedInteger. Every time we call hasNext(), we push into all NestedInteger if the current item is a list; otherwise (an integer), we stop and it is ready for calling next().

In other words, you can say that hasNext() always keep the top item of stack is an integer. It is obvious that hasNext() returns false if the stack is empty.

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
33
public class NestedIterator implements Iterator<Integer> {
  
  private 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.size() > 0) {
      NestedInteger item = stack.peek();
      // integer
      if (item.isInteger()) {
        return true;
      }
      // list
      stack.pop();
      List<NestedInteger> list = item.getList();
      for (int i = list.size() - 1; i >= 0; --i) {
        stack.push(list.get(i));
      }
    }
    return false;
  }
}

Flattened List

Reference: Flatten the list and iterate with plain next() and hasNext() (Java)

Flatten the list in the constructor! Runtime: 2 ms, faster than 100.00%.

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 List<Integer> flattenedList;
  private Iterator<Integer> it;

  public NestedIterator(List<NestedInteger> nestedList) {
    flattenedList = new ArrayList<>();
    flatten(nestedList);
    it = flattenedList.iterator();
  }

  // recursive call
  private void flatten(List<NestedInteger> nestedList) {
    for (NestedInteger i : nestedList) {
      if (i.isInteger()) {
        flattenedList.add(i.getInteger());
      } else {
        flatten(i.getList());
      }
    }
  }

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

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