We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Flatten Nested List Iterator

Number: 341

Difficulty: Medium

Paid? No

Companies: LinkedIn, Meta, Amazon, OpenAI, Apple, Google, DE Shaw, Walmart Labs, Netflix, X, Airbnb, Yandex, Warnermedia, Tesla


Problem Description

Implement an iterator over a nested list of integers where each element could be an integer or a list of integers (nested arbitrarily). The iterator should support two methods: next(), which returns the next integer, and hasNext(), which checks if there are any more integers to iterate over. The goal is to flatten the nested structure such that successive calls to next() yield the integers in their original order.


Key Insights

  • Use a stack to manage the nested elements; pushing items in reverse order maintains the correct processing sequence.
  • In the hasNext() method, perform a lazy flattening by processing nested lists only as needed.
  • The iterator ensures that when next() is called, the top of the stack is always an integer.
  • The space complexity is linear in the total number of elements due to the auxiliary stack.

Space and Time Complexity

Time Complexity: O(n) in the worst case, where n is the total number of integers, as each element is processed once. Space Complexity: O(n) due to the use of the stack to manage the nested elements.


Solution

The solution uses a stack to store the NestedInteger elements in a reversed order. When hasNext() is invoked, the algorithm checks the top element of the stack:

  • If it is an integer, we are ready to call next().
  • If it is a list, pop it from the stack and push its elements (also in reverse order) back onto the stack. This lazy processing ensures that we flatten the nested structure on the fly and maintain O(n) time and space complexity.

Code Solutions

Below are the implementations in Python, JavaScript, C++, and Java.

# The NestedIterator class uses a stack to hold the elements from the nested list.
class NestedIterator:
    def __init__(self, nestedList):
        # Initialize the stack with the reversed nested list.
        self.stack = list(reversed(nestedList))
    
    def next(self):
        # hasNext() ensures that the top is an integer.
        # Pop and return the integer from the top of the stack.
        return self.stack.pop().getInteger()
    
    def hasNext(self):
        # Process the stack until we find an integer at the top.
        while self.stack:
            top = self.stack[-1]
            if top.isInteger():
                return True
            # Pop the list at the top and push its elements in reverse order.
            self.stack.pop()
            nested = top.getList()
            for item in reversed(nested):
                self.stack.append(item)
        return False
← Back to All Questions