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 vector vec.
  • next() returns the next element from the 2D vector and moves the pointer one step forward. You may assume that all the calls to next are valid.
  • hasNext() returns true if there are still some elements in the vector, and false otherwise.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
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 <= 200
  • 0 <= vec[i].length <= 500
  • -500 <= vec[i][j] <= 500
  • At most 105 calls will be made to next and hasNext.

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

  1. Initialize two indices: row for the current row and col for the current column.
  2. In hasNext(), advance row and col to skip empty rows or columns until a valid element is found or the end is reached.
  3. In next(), call hasNext() to ensure the pointers are at a valid element, then return the current element and increment col.
  4. Handle edge cases where rows may be empty or the entire vector is empty.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
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();
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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)
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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();
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
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
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
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)
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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) per next() and hasNext() call (amortized, since each element is visited once).
  • 🧺 Space complexity: O(1) (no extra space except pointers).