LeetCode 251 - Flatten 2D Vector

The problem asks us to design an iterator that traverses a two dimensional array as though it were a single flat sequence. Instead of returning nested lists, the iterator should expose elements one at a time through two methods, next() and hasNext().

LeetCode Problem 251

Difficulty: 🟡 Medium
Topics: Array, Two Pointers, Design, Iterator

Solution

Problem Understanding

The problem asks us to design an iterator that traverses a two dimensional array as though it were a single flat sequence. Instead of returning nested lists, the iterator should expose elements one at a time through two methods, next() and hasNext().

The input is a 2D vector, meaning a list of lists. Each inner list may contain zero or more integers. The iterator must visit elements in row-major order, meaning it processes all elements in the first row before moving to the second row, and so on.

For example, given:

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

the flattened traversal order becomes:

1, 2, 3, 4

The next() method returns the next available integer and advances the iterator forward by one position. The hasNext() method checks whether another element still exists somewhere in the remaining structure.

One important detail is that the problem guarantees all calls to next() are valid. That means we never need to handle situations where next() is called after the iterator is exhausted.

The constraints are relatively small in terms of rows, but the total number of elements can still be large:

  • Up to 200 rows
  • Each row can contain up to 500 elements
  • Up to 10^5 calls to next() and hasNext()

These constraints strongly suggest that repeatedly scanning the entire structure would be inefficient. Since hasNext() may be called many times, we need an implementation where both operations are close to constant time.

Several edge cases are important:

  • The outer vector may be empty, []
  • Some inner vectors may be empty, [[], [], [1]]
  • Consecutive empty rows may appear
  • hasNext() may be called multiple times before next()
  • The final row may be empty

A naive implementation often fails on empty inner arrays because advancing between rows is easy to overlook.

Approaches

Brute Force Approach

The simplest approach is to flatten the entire 2D vector during initialization into a single one dimensional array.

For example:

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

would immediately become:

[1,2,3,4]

The iterator would then simply maintain a single index into this flattened list.

This approach is correct because flattening preserves the original traversal order. Every call to next() returns the next element in the flattened sequence, and hasNext() checks whether the index is still within bounds.

However, this solution uses extra memory proportional to the total number of elements because it duplicates the data into another array. While acceptable for the given constraints, it is not optimal from a design perspective because the iterator could instead traverse the original structure directly.

Optimal Approach

The better approach is to maintain two pointers:

  • A row pointer
  • A column pointer

These pointers represent the current position inside the original 2D vector.

The key observation is that we do not actually need to flatten the structure ahead of time. We only need a way to skip empty rows and locate the next valid element when requested.

The iterator works lazily:

  • hasNext() advances the pointers until either:

  • a valid element is found, or

  • all rows are exhausted

  • next() simply returns the current element and advances the column pointer

This approach avoids copying data and uses only constant extra memory.

Approach Time Complexity Space Complexity Notes
Brute Force O(N) initialization, O(1) per operation O(N) Flattens entire vector into a separate array
Optimal O(1) amortized per operation O(1) Traverses original structure lazily using row and column pointers

Here, N represents the total number of elements across all rows.

Algorithm Walkthrough

Optimal Lazy Iterator Algorithm

  1. Store the original 2D vector inside the class.

We keep a reference to the input structure instead of flattening it. This avoids unnecessary copying and reduces memory usage. 2. Initialize two indices:

  • row
  • col

Both start at 0, representing the first element position. 3. Implement hasNext() so it advances past empty rows.

While:

  • row is still inside the vector, and
  • col has reached the end of the current row

move to the next row and reset col back to 0.

This step is crucial because some rows may be empty. 4. After skipping exhausted rows, determine whether a valid position exists.

If row < len(vec), another element exists.

Otherwise, the iterator is exhausted. 5. Implement next() by first ensuring a valid position exists.

Since the problem guarantees valid calls, we simply:

  • read vec[row][col]
  • increment col
  • return the value
  1. Future calls to hasNext() continue advancing as needed.

The iterator always progresses forward and never revisits earlier positions.

Why it works

The algorithm maintains the invariant that after every successful hasNext() call, the (row, col) pair either points to a valid element or the iterator is completely exhausted.

Because rows are processed sequentially and columns advance left to right, every element is visited exactly once in the correct order. Empty rows are skipped safely, ensuring no invalid indexing occurs.

Python Solution

from typing import List

class Vector2D:

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

    def next(self) -> int:
        value = self.vec[self.row][self.col]
        self.col += 1
        return value

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

        return self.row < len(self.vec)

# Your Vector2D object will be instantiated and called as such:
# obj = Vector2D(vec)
# param_1 = obj.next()
# param_2 = obj.hasNext()

The implementation stores the original 2D vector and maintains two pointers, row and col.

The hasNext() method is the core of the solution. Its responsibility is to move the iterator forward until either:

  • a valid element position is found, or
  • all rows are exhausted

The loop:

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

checks whether the current row has been fully consumed. If so, the iterator moves to the next row and resets the column index.

The next() method is intentionally simple. Because the problem guarantees valid calls, we can safely access the current position directly. After returning the current value, the column pointer advances by one.

An important detail is that next() does not itself skip empty rows. That logic remains centralized inside hasNext(), keeping responsibilities clean and preventing duplicated logic.

Go Solution

type Vector2D struct {
	vec [][]int
	row int
	col int
}

func Constructor(vec [][]int) Vector2D {
	return Vector2D{
		vec: vec,
		row: 0,
		col: 0,
	}
}

func (this *Vector2D) Next() int {
	value := this.vec[this.row][this.col]
	this.col++
	return value
}

func (this *Vector2D) HasNext() bool {
	for this.row < len(this.vec) && this.col >= len(this.vec[this.row]) {
		this.row++
		this.col = 0
	}

	return this.row < len(this.vec)
}

/**
 * Your Vector2D object will be instantiated and called as such:
 * obj := Constructor(vec);
 * param_1 := obj.Next();
 * param_2 := obj.HasNext();
 */

The Go implementation closely mirrors the Python version. The main difference is that Go uses a struct with explicit fields instead of instance attributes.

Slices in Go naturally support empty inner arrays, so no special handling is required for nil or empty rows. The same pointer advancement logic works correctly.

Integer overflow is not a concern because indices remain well within Go's integer range.

Worked Examples

Example 1

Input:

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

Initial state:

Step Operation row col Returned
1 initialize 0 0 -

Call next():

Step Action row col Returned
2 read vec[0][0] 0 1 1

Call next() again:

Step Action row col Returned
3 read vec[0][1] 0 2 2

Call next() again:

Before returning, hasNext() logic would move to the next valid row because col == len(vec[0]).

Step Action row col Returned
4 move to next row 1 0 -
5 read vec[1][0] 1 1 3

Call hasNext():

Step Action row col Result
6 move to row 2 2 0 true

Call next():

Step Action row col Returned
7 read vec[2][0] 2 1 4

Final hasNext():

Step Action row col Result
8 move beyond final row 3 0 false

Example with Empty Rows

Input:

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

Initial state:

Step Action row col
1 initialize 0 0

Calling hasNext():

Step Action row col
2 skip empty row 0 1 0
3 skip empty row 1 2 0
4 valid row found 2 0

The iterator correctly skips empty rows and positions itself at value 1.

Complexity Analysis

Measure Complexity Explanation
Time O(1) amortized Each element and row is processed at most once
Space O(1) Only two indices are stored

Although hasNext() may occasionally skip multiple rows, each row transition happens only once during the entire lifetime of the iterator. Therefore, the total work across all operations is linear in the number of rows and elements, giving amortized constant time per operation.

Test Cases

from typing import List

class Vector2D:

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

    def next(self) -> int:
        value = self.vec[self.row][self.col]
        self.col += 1
        return value

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

        return self.row < len(self.vec)

# Provided example
obj = Vector2D([[1, 2], [3], [4]])
assert obj.next() == 1           # first element
assert obj.next() == 2           # second element
assert obj.next() == 3           # move across rows
assert obj.hasNext() is True     # still has 4
assert obj.hasNext() is True     # repeated hasNext calls
assert obj.next() == 4           # final element
assert obj.hasNext() is False    # exhausted iterator

# Empty outer vector
obj = Vector2D([])
assert obj.hasNext() is False    # completely empty input

# Only empty rows
obj = Vector2D([[], [], []])
assert obj.hasNext() is False    # no actual elements

# Empty rows between valid rows
obj = Vector2D([[1], [], [2, 3]])
assert obj.hasNext() is True
assert obj.next() == 1
assert obj.hasNext() is True
assert obj.next() == 2
assert obj.next() == 3
assert obj.hasNext() is False

# Single element
obj = Vector2D([[42]])
assert obj.hasNext() is True
assert obj.next() == 42
assert obj.hasNext() is False

# Multiple repeated hasNext calls
obj = Vector2D([[5, 6]])
assert obj.hasNext() is True
assert obj.hasNext() is True
assert obj.hasNext() is True
assert obj.next() == 5
assert obj.next() == 6
assert obj.hasNext() is False

# Large empty prefix
obj = Vector2D([[], [], [], [7]])
assert obj.hasNext() is True
assert obj.next() == 7
assert obj.hasNext() is False
Test Why
[[1,2],[3],[4]] Validates standard traversal
[] Tests completely empty input
[[],[],[]] Ensures empty rows are skipped correctly
[[1],[],[2,3]] Tests gaps between populated rows
[[42]] Smallest non-empty case
repeated hasNext() calls Ensures state is not corrupted
[[],[],[],[7]] Tests long sequences of empty rows

Edge Cases

Empty Outer Vector

An input like:

[]

contains no rows at all. A naive implementation might immediately attempt to access vec[0], causing an index error.

This implementation handles the case safely because hasNext() first checks:

self.row < len(self.vec)

If the vector is empty, the condition fails immediately and returns False.

Empty Inner Rows

Inputs such as:

[[], [], [1]]

are particularly tricky because the iterator must skip invalid rows before reaching a valid element.

A buggy implementation may incorrectly assume every row contains at least one value. This solution avoids that issue by repeatedly advancing rows inside the while loop in hasNext() until a row with remaining elements is found.

Repeated Calls to hasNext()

Many iterator problems fail when hasNext() is called multiple times consecutively.

For example:

obj.hasNext()
obj.hasNext()
obj.hasNext()

should not advance the iterator multiple times or skip elements.

This implementation works correctly because hasNext() only moves forward when the current row is exhausted. Once positioned at a valid element, repeated calls leave the state unchanged.