LeetCode/java/flatten-nested-list-iterator.java

import java.util.ArrayDeque;
import java.util.Iterator;
import java.util.List;

// This is the interface that allows for creating nested lists.
// You should not implement it, or speculate about its implementation
interface NestedInteger {
  // @return true if this NestedInteger holds a single integer, rather than a nested list.
  public boolean isInteger();

  // @return the single integer that this NestedInteger holds, if it holds a single integer
  // Return null if this NestedInteger holds a nested list
  public Integer getInteger();

  // @return the nested list that this NestedInteger holds, if it holds a nested list
  // Return empty list if this NestedInteger holds a single integer
  public List<NestedInteger> getList();
}

class NestedIterator implements Iterator<Integer> {

  private class DFSEntry {
    private List<NestedInteger> it;
    private int i = 0;

    public DFSEntry(List<NestedInteger> it) {
      this.it = it;
    }

    private int size() {
      return it.size();
    }

    public boolean isEmpty() {
      return i >= size();
    }

    public NestedInteger at() {
      return it.get(i);
    }

    public void advance() {
      i++;
    }
  }

  ArrayDeque<DFSEntry> stack = new ArrayDeque<>();

  public NestedIterator(List<NestedInteger> nestedList) {
    stack.addLast(new DFSEntry(nestedList));
    findNext();
  }

  private DFSEntry top() {
    return stack.peekLast();
  }

  // Performs one „iteration“ of the DFS algorithm
  private void findNext() {
    while (!stack.isEmpty()) {
      // We ran out of successors
      if (top().isEmpty()) {
        stack.removeLast();

        // Advance the parent of the removed list
        if (top() != null) {
          top().advance();
        }

        continue;
      }

      // Base: we have found an integer
      if (top().at().isInteger()) {
        return;
      }

      stack.addLast(new DFSEntry(top().at().getList()));
    }
  }

  @Override
  public Integer next() {
    DFSEntry top = stack.getLast();
    Integer e = top.at().getInteger();

    // Find next integer in the nested list
    top.advance();
    findNext();

    // Return the popped element
    return e;
  }

  @Override
  public boolean hasNext() {
    return !stack.isEmpty();
  }
}

/**
 * Your NestedIterator object will be instantiated and called as such: NestedIterator i = new
 * NestedIterator(nestedList); while (i.hasNext()) v[f()] = i.next();
 */