algoadvance

You are given a nested list of integers. Each element is either an integer or a list whose elements may also be integers or other lists. Implement an iterator to flatten it.

Implement the NestedIterator class:

Your code will be tested with the following pseudocode:

initialize the nestedList with a NestedIterator
repeat the following while nestedList.hasNext():
    print(nestedList.next())

This will generate output of all the integers in the nested list in the same order they appear in the nested list.

Clarifying Questions

  1. What is the structure of NestedInteger?
    • NestedInteger is a class with the following methods:
      • isInteger(): Returns true if this NestedInteger holds a single integer, rather than a nested list.
      • getInteger(): Returns the single integer that this NestedInteger holds, if it holds a single integer. The result is undefined if this NestedInteger holds a nested list.
      • getList(): Returns the nested list that this NestedInteger holds, if it holds a nested list. The result is undefined if this NestedInteger holds a single integer.
  2. Can the input list be empty?
    • Yes, the input list can be empty, in which case hasNext() should return false.

Strategy

The goal is to flatten out a nested list structure such that we can iterate over all the integers within. A useful approach is to use a stack to maintain the position within the nested lists.

Here’s the overall approach:

  1. Initialization (__init__):
    • We’ll use a stack to store iterators of lists, starting with the main list.
    • We need to ensure that we only keep non-empty iterators on the stack, hence some preprocessing might be required to clean up the stack.
  2. Next Integer (next):
    • If hasNext() confirms that there is a next integer, simply return it.
  3. Check for Next Integer (hasNext):
    • Process the stack to find the next integer.
    • If the current iterator points to a NestedInteger, check if it is an integer or a list:
      • If it’s an integer, keep it as the next integer.
      • If it’s a list, push its iterator onto the stack and keep processing.
    • Clean up invalid/empty iterators from the stack.

Code

class NestedIterator:
    def __init__(self, nestedList):
        self.stack = [iter(nestedList)]
        self.next_element = None

    def next(self):
        if not self.hasNext():
            raise StopIteration
        result = self.next_element
        self.next_element = None
        return result

    def hasNext(self):
        while self.stack:
            try:
                current = next(self.stack[-1])
            except StopIteration:
                self.stack.pop()
                continue
            
            if current.isInteger():
                self.next_element = current.getInteger()
                return True
            else:
                self.stack.append(iter(current.getList()))
        return False

Time Complexity

Thus, the amortized time complexity for both next and hasNext is O(1) per call.

Try our interview co-pilot at AlgoAdvance.com