Learning On the Way

341. Flatten Nested List Iterator

You are given a nested list of integers nestedList. 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:

NestedIterator(List nestedList) Initializes the iterator with the nested list nestedList. int next() Returns the next integer in the nested list. boolean hasNext() Returns true if there are still some integers in the nested list and false otherwise.

Stack solution

class NestedIterator: 
     
    def __init__(self, nestedList: [NestedInteger]): 
        self.stack = list(reversed(nestedList)) 
         
         
    def next(self) -> int: 
        self.make_stack_top_an_integer() 
        return self.stack.pop().getInteger() 
     
         
    def hasNext(self) -> bool: 
        self.make_stack_top_an_integer() 
        return len(self.stack) > 0 
         
         
    def make_stack_top_an_integer(self): 
        # While the stack contains a nested list at the top... 
        while self.stack and not self.stack[-1].isInteger(): 
            # Unpack the list at the top by putting its items onto 
            # the stack in reverse order. 
            self.stack.extend(reversed(self.stack.pop().getList()))

Recursive solution:

class NestedIterator: 
     
    def __init__(self, nestedList: [NestedInteger]): 
        def flatten_list(nested_list): 
            for nested_integer in nested_list: 
                if nested_integer.isInteger(): 
                    self._integers.append(nested_integer.getInteger()) 
                else: 
                    flatten_list(nested_integer.getList())  
        self._integers = [] 
        self._position = -1 # Pointer to previous returned. 
        flatten_list(nestedList) 
     
    def next(self) -> int: 
        self._position += 1 
        return self._integers[self._position] 
         
    def hasNext(self) -> bool: 
        return self._position + 1 < len(self._integers)

Two stacks to keep index:

class NestedIterator: 
     
    def __init__(self, nestedList: [NestedInteger]): 
        self.stack = [[nestedList, 0]] 
         
    def make_stack_top_an_integer(self): 
         
        while self.stack: 
             
            # Essential for readability :) 
            current_list = self.stack[-1][0] 
            current_index = self.stack[-1][1] 
             
            # If the top list is used up, pop it and its index. 
            if len(current_list) == current_index: 
                self.stack.pop() 
                continue 
             
            # Otherwise, if it's already an integer, we don't need  
            # to do anything. 
            if current_list[current_index].isInteger(): 
                break 
             
            # Otherwise, it must be a list. We need to increment the index 
            # on the previous list, and add the new list. 
            new_list = current_list[current_index].getList() 
            self.stack[-1][1] += 1 # Increment old. 
            self.stack.append([new_list, 0]) 
             
     
    def next(self) -> int: 
        self.make_stack_top_an_integer() 
        current_list = self.stack[-1][0] 
        current_index = self.stack[-1][1] 
        self.stack[-1][1] += 1 
        return current_list[current_index].getInteger() 
         
     
    def hasNext(self) -> bool: 
        self.make_stack_top_an_integer() 
        return len(self.stack) > 0