LeetCode 251: Flatten 2D Vector

Problem Restatement

We need to design an iterator over a 2D vector.

The input is a list of lists, for example:

vec = [[1, 2], [3], [4]]

The iterator should behave as if the numbers were stored in one flat list:

[1, 2, 3, 4]

The class must support two operations:

The problem guarantees that next() is called only when a next element exists. The 2D vector may contain empty inner lists. At most 10^5 calls are made to next() and hasNext().

Input and Output

Example class shape:

class Vector2D:

    def __init__(self, vec: list[list[int]]):
        ...

    def next(self) -> int:
        ...

    def hasNext(self) -> bool:
        ...

Examples

Consider:

vec = [[1, 2], [3], [4]]

Calls:

iterator = Vector2D(vec)

iterator.next()     # 1
iterator.next()     # 2
iterator.next()     # 3
iterator.hasNext()  # True
iterator.hasNext()  # True
iterator.next()     # 4
iterator.hasNext()  # False

The flattened order is:

[1, 2, 3, 4]

Now consider empty rows:

vec = [[], [1], [], [2, 3], []]

The iterator should skip empty lists and return:

1, 2, 3

First Thought: Flatten Everything

The simplest solution is to flatten the entire 2D vector in the constructor.

class Vector2D:

    def __init__(self, vec: list[list[int]]):
        self.data = []

        for row in vec:
            for value in row:
                self.data.append(value)

        self.index = 0

    def next(self) -> int:
        value = self.data[self.index]
        self.index += 1
        return value

    def hasNext(self) -> bool:
        return self.index < len(self.data)

This is easy to understand.

The constructor builds:

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

into:

[1, 2, 3, 4]

Then next() only reads from a normal list.

Problem With Flattening

Flattening uses extra memory.

If the input contains many values, we duplicate all of them into another list. That gives us:

Here, n is the total number of integers in the 2D vector.

We can do better in space by keeping pointers into the original vector.

Key Insight

A 2D vector has two coordinates:

vec[row][col]

So instead of creating a flat list, we store two pointers:

For example:

vec = [[1, 2], [3], [4]]

At the beginning:

row = 0
col = 0

So the current value is:

vec[0][0] = 1

After returning 1, we move to:

row = 0
col = 1

Now the current value is:

vec[0][1] = 2

After returning 2, col moves past the end of row 0, so we move to the next row:

row = 1
col = 0

Now the current value is:

vec[1][0] = 3

The main detail is skipping empty rows.

Algorithm

Store:

self.vec = vec
self.row = 0
self.col = 0

Before answering hasNext() or reading the next value, move the pointers until either:

1. row points to a row with an unread element, or
2. row reaches the end of vec.

We can put that logic into a helper method called _advance.

def _advance(self) -> None:
    while self.row < len(self.vec) and self.col == len(self.vec[self.row]):
        self.row += 1
        self.col = 0

This loop skips exhausted rows and empty rows.

Then:

hasNext()

calls _advance() and checks whether row is still inside the vector.

next()

also calls _advance(), reads vec[row][col], then increments col.

Correctness

The iterator keeps the invariant that after _advance() finishes, one of two conditions holds.

Either row == len(vec), which means there are no rows left, so no next element exists.

Or row < len(vec) and col < len(vec[row]), which means vec[row][col] is the next unread integer in row-major order.

The helper method only moves forward. It never skips a valid unread integer, because it advances to the next row only when col == len(vec[row]), meaning the current row has already been fully consumed or is empty.

When next() returns vec[row][col], it returns exactly the next unread integer. Then it increments col, so the same integer cannot be returned again.

Therefore, repeated calls to next() return all integers from the 2D vector in left-to-right, top-to-bottom order.

Complexity

Let n be the total number of integers and r be the number of rows.

A single hasNext() call may skip several empty rows, but across all calls each row is skipped at most once. So the amortized cost is constant.

Implementation

class Vector2D:

    def __init__(self, vec: list[list[int]]):
        self.vec = vec
        self.row = 0
        self.col = 0

    def _advance(self) -> None:
        while self.row < len(self.vec) and self.col == len(self.vec[self.row]):
            self.row += 1
            self.col = 0

    def next(self) -> int:
        self._advance()

        value = self.vec[self.row][self.col]
        self.col += 1

        return value

    def hasNext(self) -> bool:
        self._advance()
        return self.row < len(self.vec)

Code Explanation

The constructor stores the original vector:

self.vec = vec

It also starts the iterator at the first possible position:

self.row = 0
self.col = 0

The helper method skips rows that cannot produce a value:

while self.row < len(self.vec) and self.col == len(self.vec[self.row]):

This condition handles two cases.

For an empty row:

[]

we have:

col == len(row) == 0

So we move past it.

For an exhausted row:

[1, 2]

after returning both elements, col == 2, which equals the row length. So we move to the next row.

The next() method first normalizes the pointer position:

self._advance()

Then it reads the current value:

value = self.vec[self.row][self.col]

After reading, it moves one step to the right:

self.col += 1

The hasNext() method also calls _advance().

self._advance()
return self.row < len(self.vec)

If row is still inside the vector, there is another value.

If row has reached the end, the iterator is exhausted.

Testing

def collect(vec: list[list[int]]) -> list[int]:
    iterator = Vector2D(vec)
    result = []

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

    return result

def run_tests():
    assert collect([[1, 2], [3], [4]]) == [1, 2, 3, 4]
    assert collect([[], [1], [], [2, 3], []]) == [1, 2, 3]
    assert collect([]) == []
    assert collect([[], [], []]) == []
    assert collect([[1]]) == [1]
    assert collect([[1, 2, 3], [], [4], [5, 6]]) == [1, 2, 3, 4, 5, 6]

    iterator = Vector2D([[1], [], [2]])
    assert iterator.hasNext() is True
    assert iterator.hasNext() is True
    assert iterator.next() == 1
    assert iterator.hasNext() is True
    assert iterator.next() == 2
    assert iterator.hasNext() is False

    print("all tests passed")

run_tests()

Test meaning: