Flatten 2D Vector
MediumUpdated: Aug 2, 2025
Practice on:
Problem
Design an iterator to flatten a 2D vector. It should support the next and hasNext operations.
Implement the Vector2D class:
Vector2D(int[][] vec)initializes the object with the 2D vectorvec.next()returns the next element from the 2D vector and moves the pointer one step forward. You may assume that all the calls tonextare valid.hasNext()returnstrueif there are still some elements in the vector, andfalseotherwise.
Examples
Example 1
Input
["Vector2D", "next", "next", "next", "hasNext", "hasNext", "next", "hasNext"]
[[[[1, 2], [3], [4]]], [], [], [], [], [], [], []]
Output
[null, 1, 2, 3, true, true, 4, false]
Explanation
Vector2D vector2D = new Vector2D([[1, 2], [3], [4]]);
vector2D.next(); // return 1
vector2D.next(); // return 2
vector2D.next(); // return 3
vector2D.hasNext(); // return True
vector2D.hasNext(); // return True
vector2D.next(); // return 4
vector2D.hasNext(); // return False
Constraints
0 <= vec.length <= 2000 <= vec[i].length <= 500-500 <= vec[i][j] <= 500- At most
105calls will be made tonextandhasNext.
Follow up
As an added challenge, try to code it using only iterators in C++ or iterators in Java.
Solution
Method 1 – Two Pointers Iterator
Intuition
The key idea is to use two pointers to track the current row and column in the 2D vector. We advance the pointers to skip empty rows and always point to the next available element. This allows us to flatten the 2D vector on-the-fly without extra space.
Approach
- Initialize two indices:
rowfor the current row andcolfor the current column. - In
hasNext(), advancerowandcolto skip empty rows or columns until a valid element is found or the end is reached. - In
next(), callhasNext()to ensure the pointers are at a valid element, then return the current element and incrementcol. - Handle edge cases where rows may be empty or the entire vector is empty.
Code
C++
class Vector2D {
vector<vector<int>> &vec;
int row, col;
public:
Vector2D(vector<vector<int>>& v) : vec(v), row(0), col(0) {}
int next() {
hasNext();
return vec[row][col++];
}
bool hasNext() {
while (row < vec.size() && col >= vec[row].size()) {
row++;
col = 0;
}
return row < vec.size();
}
};
Go
type Vector2D struct {
vec [][]int
row, col int
}
func Constructor(vec [][]int) Vector2D {
return Vector2D{vec: vec}
}
func (v *Vector2D) Next() int {
v.HasNext()
ans := v.vec[v.row][v.col]
v.col++
return ans
}
func (v *Vector2D) HasNext() bool {
for v.row < len(v.vec) && v.col >= len(v.vec[v.row]) {
v.row++
v.col = 0
}
return v.row < len(v.vec)
}
Java
class Vector2D {
private List<List<Integer>> vec;
private int row, col;
public Vector2D(int[][] v) {
vec = new ArrayList<>();
for (int[] arr : v) {
List<Integer> list = new ArrayList<>();
for (int num : arr) list.add(num);
vec.add(list);
}
row = 0;
col = 0;
}
public int next() {
hasNext();
return vec.get(row).get(col++);
}
public boolean hasNext() {
while (row < vec.size() && col >= vec.get(row).size()) {
row++;
col = 0;
}
return row < vec.size();
}
}
Kotlin
class Vector2D(v: Array<IntArray>) {
private val vec = v
private var row = 0
private var col = 0
fun next(): Int {
hasNext()
return vec[row][col++]
}
fun hasNext(): Boolean {
while (row < vec.size && col >= vec[row].size) {
row++
col = 0
}
return row < vec.size
}
}
Python
class Vector2D:
def __init__(self, vec: list[list[int]]) -> None:
self.vec = vec
self.row = 0
self.col = 0
def next(self) -> int:
self.hasNext()
ans = self.vec[self.row][self.col]
self.col += 1
return ans
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)
Rust
struct Vector2D {
vec: Vec<Vec<i32>>,
row: usize,
col: usize,
}
impl Vector2D {
fn new(vec: Vec<Vec<i32>>) -> Self {
Vector2D { vec, row: 0, col: 0 }
}
fn next(&mut self) -> i32 {
self.has_next();
let ans = self.vec[self.row][self.col];
self.col += 1;
ans
}
fn has_next(&mut self) -> bool {
while self.row < self.vec.len() && self.col >= self.vec[self.row].len() {
self.row += 1;
self.col = 0;
}
self.row < self.vec.len()
}
}
Complexity
- ⏰ Time complexity:
O(1)pernext()andhasNext()call (amortized, since each element is visited once). - 🧺 Space complexity:
O(1)(no extra space except pointers).