341. Flatten Nested List Iterator

Approach 1: Recursive (Queue)

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class NestedIterator {
 public:
  NestedIterator(vector<NestedInteger>& nestedList) {
    addInteger(nestedList);
  }

  int next() {
    const int num = q.front();
    q.pop();
    return num;
  }

  bool hasNext() {
    return !q.empty();
  }

 private:
  queue<int> q;

  void addInteger(const vector<NestedInteger>& nestedList) {
    for (const NestedInteger& ni : nestedList)
      if (ni.isInteger())
        q.push(ni.getInteger());
      else
        addInteger(ni.getList());
  }
};

public class NestedIterator implements Iterator<Integer> {
  public NestedIterator(List<NestedInteger> nestedList) {
    addInteger(nestedList);
  }

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

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

  private Queue<Integer> q = new ArrayDeque<>();

  private void addInteger(final List<NestedInteger> nestedList) {
    for (final NestedInteger ni : nestedList)
      if (ni.isInteger())
        q.offer(ni.getInteger());
      else
        addInteger(ni.getList());
  }
}

class NestedIterator:
  def __init__(self, nestedList: list[NestedInteger]):
    self.q = collections.deque()
    self.addInteger(nestedList)

  def next(self) -> int:
    return self.q.popleft()

  def hasNext(self) -> bool:
    return self.q

  def addInteger(self, nestedList: list[NestedInteger]) -> None:
    for ni in nestedList:
      if ni.isInteger():
        self.q.append(ni.getInteger())
      else:
        self.addInteger(ni.getList())

Approach 2: Stack

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class NestedIterator {
 public:
  NestedIterator(vector<NestedInteger>& nestedList) {
    addInteger(nestedList);
  }

  int next() {
    const int num = stack.top().getInteger();
    stack.pop();
    return num;
  }

  bool hasNext() {
    while (!stack.empty() && !stack.top().isInteger()) {
      vector<NestedInteger> nestedList = stack.top().getList();
      stack.pop();
      addInteger(nestedList);
    }
    return !stack.empty();
  }

 private:
  stack<NestedInteger> stack;

  // addInteger([1, [4, [6]]]) -> stack = [[4, [6]], 1]
  // addInteger([4, [6]]) -> stack = [[6], 4]
  // addInteger([6]) -> stack = [6]
  void addInteger(const vector<NestedInteger>& nestedList) {
    for (int i = nestedList.size() - 1; i >= 0; --i)
      stack.push(nestedList[i]);
  }
};

public class NestedIterator implements Iterator<Integer> {
  public NestedIterator(List<NestedInteger> nestedList) {
    addInteger(nestedList);
  }

  @Override
  public Integer next() {
    return stack.pop().getInteger();
  }

  @Override
  public boolean hasNext() {
    while (!stack.isEmpty() && !stack.peek().isInteger()) {
      final NestedInteger ni = stack.pop();
      addInteger(ni.getList());
    }
    return !stack.isEmpty();
  }

  private Deque<NestedInteger> stack = new ArrayDeque<>();

  // addInteger([1, [4, [6]]]) -> stack = [[4, [6]], 1]
  // addInteger([4, [6]]) -> stack = [[6], 4]
  // addInteger([6]) -> stack = [6]
  private void addInteger(final List<NestedInteger> nestedList) {
    for (int i = nestedList.size() - 1; i >= 0; --i)
      stack.push(nestedList.get(i));
  }
}

class NestedIterator:
  def __init__(self, nestedList: list[NestedInteger]):
    self.stack: list[NestedInteger] = []
    self.addInteger(nestedList)

  def next(self) -> int:
    return self.stack.pop().getInteger()

  def hasNext(self) -> bool:
    while self.stack and not self.stack[-1].isInteger():
      self.addInteger(self.stack.pop().getList())
    return self.stack

  # addInteger([1, [4, [6]]]) . stack = [[4, [6]], 1]
  # addInteger([4, [6]]) . stack = [[6], 4]
  # addInteger([6]) . stack = [6]
  def addInteger(self, nestedList: list[NestedInteger]) -> None:
    for n in reversed(nestedList):
      self.stack.append(n)