Flatten a Nested List¶
Nested structures appear constantly, JSON responses, recursive configs, tree-structured data. This problem asks you to iterate over an arbitrarily nested list and yield elements one at a time, flat.
Problem¶
Design an iterator for a heterogeneous list where each element is either an int or a nested list of ints (to arbitrary depth).
[9, [1, 2], 10, 20] -> [9, 1, 2, 10, 20]
[9, [1, 2, 4, [44, 55, 66]], 10, 20] -> [9, 1, 2, 4, 44, 55, 66, 10, 20]
Approach¶
Two classic strategies, both doing depth-first traversal:
Stack-based (iterative): Push elements onto a stack in reverse order. On each __next__ call, pop the top, if it's a list, expand it onto the stack (reversed); if it's an int, return it. This avoids deep call stacks, which matters when nesting depth is unbounded.
Recursive: Walk each element; if it's a list, recurse into it; if it's an int, append to the output. Simpler to read, but risks stack overflow on deeply nested input.
Implementation¶
Stack-based iterator¶
class FlattenListIterator:
def __init__(self, nested_list):
self.stack = list(reversed(nested_list))
def __iter__(self):
return self
def __next__(self):
while self.stack:
top = self.stack.pop()
if isinstance(top, list):
self.stack.extend(reversed(top))
else:
return top
raise StopIteration
nested = [9, [1, 2, 4, [44, 55, 66]], 10, 20]
print(list(FlattenListIterator(nested)))
# [9, 1, 2, 4, 44, 55, 66, 10, 20]
Recursive¶
def flatten(nested_list):
result = []
for item in nested_list:
if isinstance(item, list):
result.extend(flatten(item))
else:
result.append(item)
return result
print(flatten([9, [1, 2], 10, 20]))
# [9, 1, 2, 10, 20]
Complexity¶
- Time: O(n) where n is the total number of integers across all nesting levels
- Space: O(n) for the stack / recursion depth in the worst case
Takeaway¶
The stack-based iterator is the production-safe choice, it handles arbitrary nesting without blowing the call stack and yields elements lazily, which is useful when you're streaming data through a pipeline and don't want to materialize the full flat list in memory.