LeetCode 341 - Flatten Nested List Iterator

This problem asks us to design an iterator that can traverse a deeply nested list structure as if it were a flat list of integers. The input is not a normal array.

leetcodemediumstacktreedepth-first-searchdesignqueueiterator

LeetCode Problem 341

Difficulty: 🟡 Medium
Topics: Stack, Tree, Depth-First Search, Design, Queue, Iterator

Solution

Problem Understanding

This problem asks us to design an iterator that can traverse a deeply nested list structure as if it were a flat list of integers.

The input is not a normal array. Each element inside nestedList can be one of two things:

  • A single integer
  • Another list containing more NestedInteger objects

Those nested lists may themselves contain additional nested lists, potentially many levels deep.

The goal is to implement an iterator with the standard iterator interface:

  • next() returns the next integer in flattened order
  • hasNext() tells us whether another integer still exists

The iterator should behave exactly as though the entire nested structure had been flattened into a simple one-dimensional array from left to right.

For example:

[[1,1],2,[1,1]]

should behave like:

[1,1,2,1,1]

and:

[1,[4,[6]]]

should behave like:

[1,4,6]

The important detail is that flattening must preserve the original left-to-right traversal order.

The constraints are moderate:

  • Up to 500 top-level elements
  • Arbitrary nesting depth
  • Integer values can range from -10^6 to 10^6

The relatively small input size means even an eager flattening approach is acceptable. However, the problem is fundamentally about iterator design, so a lazy traversal solution is usually considered the more elegant and interview-preferred approach.

Several edge cases are important:

  • A list may contain only integers
  • A list may contain only nested empty lists
  • Nesting may be extremely deep
  • Lists may alternate between integers and sublists
  • The first element may itself be deeply nested

A naive implementation can easily fail if it does not properly preserve traversal order or if it incorrectly handles nested empty lists.

Approaches

Brute Force Approach

The brute force solution is to completely flatten the nested structure during initialization.

We perform a depth-first traversal of the entire nested list and append every integer into a normal array. Then:

  • next() simply returns the next element from the flattened array
  • hasNext() checks whether the index is still within bounds

This approach is straightforward and easy to reason about because the iterator logic becomes trivial after preprocessing.

The recursive flattening process works because depth-first traversal naturally preserves the required left-to-right ordering.

However, this solution eagerly processes the entire structure up front, even if the caller only requests a few elements. It also stores all integers simultaneously in memory.

Optimal Approach

A better solution is to flatten lazily using a stack.

Instead of preprocessing everything, we only expand nested lists when necessary. The iterator maintains a stack of elements that still need to be processed.

The key insight is that stacks naturally simulate recursive depth-first traversal.

To preserve left-to-right order, whenever we encounter a nested list we push its contents onto the stack in reverse order. That ensures the leftmost element appears on top of the stack and is processed first.

This design gives us an iterator that processes elements incrementally and avoids unnecessary upfront work.

Approach Time Complexity Space Complexity Notes
Brute Force O(n) O(n) Fully flattens the structure during initialization
Optimal O(n) O(d) average, O(n) worst case Uses lazy DFS traversal with a stack

Here:

  • n is the total number of integers and nested containers
  • d is the maximum nesting depth

Algorithm Walkthrough

Lazy Stack-Based DFS

  1. Initialize a stack with all top-level elements in reverse order.

We reverse the order because stacks are LIFO structures. Reversing ensures the leftmost element is processed first. 2. Implement hasNext() so that it guarantees the top of the stack is an integer.

This is the most important part of the algorithm. 3. While the stack is not empty:

  • Look at the top element without removing it.
  • If it is an integer, we are ready to return True.
  • If it is a nested list, pop it from the stack and expand its contents.
  1. When expanding a nested list:
  • Retrieve the contained list
  • Push its elements onto the stack in reverse order

Again, reversing preserves the correct traversal order. 5. If the stack eventually becomes empty, return False.

That means no integers remain. 6. Implement next() by first calling hasNext().

This guarantees the top element is an integer before removal. 7. Pop the integer from the stack and return its value.

Why it works

The stack always represents the remaining portion of the DFS traversal.

The invariant is:

  • If hasNext() returns True, the top of the stack is guaranteed to be an integer.
  • Elements are pushed in reverse order, so popping follows the original left-to-right traversal order.

Because every nested list is expanded exactly once and every integer is returned exactly once, the iterator produces the correct flattened sequence.

Python Solution

# """
# This is the interface that allows for creating nested lists.
# You should not implement it, or speculate about its implementation
# """
# class NestedInteger:
#     def isInteger(self) -> bool:
#         pass
#
#     def getInteger(self) -> int:
#         pass
#
#     def getList(self) -> [NestedInteger]:
#         pass

from typing import List

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

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

    def hasNext(self) -> bool:
        while self.stack:
            top = self.stack[-1]

            if top.isInteger():
                return True

            self.stack.pop()

            nested = top.getList()

            for item in reversed(nested):
                self.stack.append(item)

        return False

# Your NestedIterator object will be instantiated and called as such:
# i, v = NestedIterator(nestedList), []
# while i.hasNext():
#     v.append(i.next())

The implementation uses a single stack to simulate recursive DFS traversal.

During initialization, we reverse the top-level list and store it in the stack. This ensures the first element appears on top.

The hasNext() method performs all flattening work lazily. It repeatedly examines the top element:

  • If it is already an integer, iteration can continue immediately.
  • If it is a list, the list is expanded and its contents are pushed back onto the stack in reverse order.

This guarantees that when hasNext() finishes successfully, the top stack element is always an integer.

The next() method is intentionally simple. It first ensures a valid integer exists by calling hasNext(), then pops and returns that integer.

This separation of responsibilities keeps the iterator logic clean and robust.

Go Solution

/**
 * // This is the interface that allows for creating nested lists.
 * // You should not implement it, or speculate about its implementation
 * type NestedInteger struct {
 * }
 *
 * func (this NestedInteger) IsInteger() bool {}
 * func (this NestedInteger) GetInteger() int {}
 * func (this NestedInteger) GetList() []*NestedInteger {}
 */

type NestedIterator struct {
	stack []*NestedInteger
}

func Constructor(nestedList []*NestedInteger) *NestedIterator {
	stack := make([]*NestedInteger, 0)

	for i := len(nestedList) - 1; i >= 0; i-- {
		stack = append(stack, nestedList[i])
	}

	return &NestedIterator{
		stack: stack,
	}
}

func (this *NestedIterator) Next() int {
	this.HasNext()

	n := len(this.stack)
	val := this.stack[n-1].GetInteger()

	this.stack = this.stack[:n-1]

	return val
}

func (this *NestedIterator) HasNext() bool {
	for len(this.stack) > 0 {
		top := this.stack[len(this.stack)-1]

		if top.IsInteger() {
			return true
		}

		this.stack = this.stack[:len(this.stack)-1]

		nested := top.GetList()

		for i := len(nested) - 1; i >= 0; i-- {
			this.stack = append(this.stack, nested[i])
		}
	}

	return false
}

The Go implementation follows the exact same logic as the Python version.

The primary difference is explicit slice manipulation:

  • Removing from the stack uses slicing
  • Appending expands the slice dynamically

Go also requires pointer handling because GetList() returns []*NestedInteger.

Unlike Python, Go does not allow negative indexing, so accessing the top element requires explicit length calculations.

Worked Examples

Example 1

Input:

[[1,1],2,[1,1]]

Initial Stack

We push elements in reverse order:

Stack Bottom → Top
[1,1]
2
[1,1]

Top of stack is the rightmost [1,1].

Iteration Trace

Operation Stack State Output
Expand [1,1] [1,1], 2, 1, 1 []
Pop 1 [1,1], 2, 1 [1]
Pop 1 [1,1], 2 [1,1]
Pop 2 [1,1] [1,1,2]
Expand [1,1] 1, 1 [1,1,2]
Pop 1 1 [1,1,2,1]
Pop 1 empty [1,1,2,1,1]

Final output:

[1,1,2,1,1]

Example 2

Input:

[1,[4,[6]]]

Initial Stack

Stack Bottom → Top
[4,[6]]
1

Iteration Trace

Operation Stack State Output
Pop 1 [4,[6]] [1]
Expand [4,[6]] [6], 4 [1]
Pop 4 [6] [1,4]
Expand [6] 6 [1,4]
Pop 6 empty [1,4,6]

Final output:

[1,4,6]

Complexity Analysis

Measure Complexity Explanation
Time O(n) Every integer and list is processed at most once
Space O(d) average, O(n) worst case Stack stores nested traversal state

The time complexity is linear because each element enters and leaves the stack exactly once.

The space complexity depends on nesting depth. In balanced structures, the stack only stores the current DFS path, which is proportional to depth d. In the worst case, such as one huge nested list expanded at once, the stack may temporarily contain all elements.

Test Cases

# Helper classes for local testing only

class NestedInteger:
    def __init__(self, value=None):
        if isinstance(value, int):
            self.integer = value
            self.list = None
        else:
            self.integer = None
            self.list = value if value is not None else []

    def isInteger(self):
        return self.integer is not None

    def getInteger(self):
        return self.integer

    def getList(self):
        return self.list

def build_nested(data):
    if isinstance(data, int):
        return NestedInteger(data)

    return NestedInteger([build_nested(x) for x in data])

def flatten(data):
    nested = [build_nested(x) for x in data]

    iterator = NestedIterator(nested)

    result = []

    while iterator.hasNext():
        result.append(iterator.next())

    return result

assert flatten([[1, 1], 2, [1, 1]]) == [1, 1, 2, 1, 1]  # provided example
assert flatten([1, [4, [6]]]) == [1, 4, 6]  # nested example
assert flatten([1, 2, 3]) == [1, 2, 3]  # flat list only
assert flatten([[[1]]]) == [1]  # deeply nested single value
assert flatten([[], [1], []]) == [1]  # empty nested lists
assert flatten([[[]]]) == []  # nested empty structure
assert flatten([[-1], [-2, [-3]]]) == [-1, -2, -3]  # negative integers
assert flatten([[1, [2]], 3, [[4]], 5]) == [1, 2, 3, 4, 5]  # mixed nesting
assert flatten([[[[[6]]]]]) == [6]  # extreme nesting depth
assert flatten([]) == []  # empty input
Test Why
[[1,1],2,[1,1]] Verifies standard flattening order
[1,[4,[6]]] Verifies recursive nested traversal
[1,2,3] Tests already-flat input
[[[1]]] Tests deep nesting
[[],[1],[]] Tests skipping empty lists
[[[]]] Tests fully empty nested structure
[[-1],[-2,[-3]]] Verifies negative integers
[[1,[2]],3,[[4]],5] Tests mixed nesting patterns
[[[[[6]]]]] Tests extreme recursive depth
[] Tests empty input

Edge Cases

Deeply Nested Single Integer

An input like:

[[[[[[1]]]]]]

can expose bugs in implementations that do not correctly recurse or repeatedly expand nested lists.

The stack-based approach handles this naturally because hasNext() continuously expands nested lists until the top element becomes an integer.

Nested Empty Lists

Inputs such as:

[[[]]]

are especially dangerous because the iterator must correctly skip empty containers without falsely reporting available integers.

The implementation handles this by repeatedly expanding lists until either:

  • An integer is found
  • The stack becomes empty

Empty lists simply contribute no new elements to the stack.

Alternating Integers and Lists

Inputs like:

[1,[2],3,[4,[5]]]

can break ordering if elements are not pushed in reverse order.

The algorithm explicitly reverses every expanded list before pushing onto the stack. This guarantees the original left-to-right traversal order is preserved exactly.