Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

251. Flatten 2D Vector - Leetcode Solution

Code Implementation

class Vector2D:
    def __init__(self, vec2d: list[list[int]]):
        self.vec2d = vec2d
        self.row = 0
        self.col = 0
        self._advance_to_next()

    def _advance_to_next(self):
        # Move to the next available integer
        while self.row < len(self.vec2d) and self.col >= len(self.vec2d[self.row]):
            self.row += 1
            self.col = 0

    def next(self) -> int:
        if not self.hasNext():
            raise Exception("No more elements")
        result = self.vec2d[self.row][self.col]
        self.col += 1
        self._advance_to_next()
        return result

    def hasNext(self) -> bool:
        self._advance_to_next()
        return self.row < len(self.vec2d)
      
class Vector2D {
private:
    vector<vector<int>> &vec2d;
    int row, col;
    void advanceToNext() {
        while (row < vec2d.size() && col >= vec2d[row].size()) {
            row++;
            col = 0;
        }
    }
public:
    Vector2D(vector<vector<int>>& vec) : vec2d(vec), row(0), col(0) {
        advanceToNext();
    }
    
    int next() {
        if (!hasNext()) throw runtime_error("No more elements");
        int result = vec2d[row][col];
        col++;
        advanceToNext();
        return result;
    }
    
    bool hasNext() {
        advanceToNext();
        return row < vec2d.size();
    }
};
      
public class Vector2D {
    private List<List<Integer>> vec2d;
    private int row, col;

    public Vector2D(List<List<Integer>> vec2d) {
        this.vec2d = vec2d;
        this.row = 0;
        this.col = 0;
        advanceToNext();
    }

    private void advanceToNext() {
        while (row < vec2d.size() && col >= vec2d.get(row).size()) {
            row++;
            col = 0;
        }
    }

    public int next() {
        if (!hasNext()) throw new NoSuchElementException();
        int result = vec2d.get(row).get(col);
        col++;
        advanceToNext();
        return result;
    }

    public boolean hasNext() {
        advanceToNext();
        return row < vec2d.size();
    }
}
      
var Vector2D = function(vec2d) {
    this.vec2d = vec2d;
    this.row = 0;
    this.col = 0;
    this._advanceToNext();
};

Vector2D.prototype._advanceToNext = function() {
    while (this.row < this.vec2d.length && this.col >= this.vec2d[this.row].length) {
        this.row++;
        this.col = 0;
    }
};

Vector2D.prototype.next = function() {
    if (!this.hasNext()) throw new Error("No more elements");
    var result = this.vec2d[this.row][this.col];
    this.col++;
    this._advanceToNext();
    return result;
};

Vector2D.prototype.hasNext = function() {
    this._advanceToNext();
    return this.row < this.vec2d.length;
};
      

Problem Description

You are given a 2D vector vec2d (a list of lists of integers). Implement an iterator to flatten this 2D vector. The iterator should support two operations:

  • next(): Returns the next element in the flattened vector.
  • hasNext(): Returns True if there are still elements to be returned, False otherwise.
The iterator should traverse the 2D vector in row-major order (left to right, top to bottom), skipping over empty rows if necessary. Each element should be returned exactly once and not reused.

Constraints:

  • There is exactly one valid way to flatten the vector.
  • All elements must be returned in order, with no repeats.
  • Some rows may be empty; these should be skipped.

Thought Process

At first glance, the problem asks us to process a 2D array as if it were a single flat array. A brute-force way would be to flatten the entire 2D array into a 1D array at initialization, then use a simple index to iterate. However, that approach uses extra space proportional to the number of elements.

To optimize, we can avoid storing a new array and instead keep track of our position within the original 2D structure. We need to remember both the current row and the current column. Each time we move forward, we check if we've reached the end of a row and, if so, move to the next non-empty row.

This approach is more space-efficient and handles empty rows naturally. It's similar to reading a book page by page, line by line: once you reach the end of a line, you move to the next, skipping any blank lines.

Solution Approach

The solution uses two pointers: one for the current row and one for the current column. Here's how it works step-by-step:

  1. Initialization: Set row and col to 0. These track our position in vec2d.
  2. Advance to Next Valid Position: Before accessing an element, check if the current row is within bounds and if the column index is valid for that row. If the current row is empty or we've reached the end of a row, move to the next row and reset col to 0. Repeat this until a valid position is found or we've reached the end.
  3. next(): Return the element at vec2d[row][col], then increment col. After incrementing, call the advance function to move to the next valid position.
  4. hasNext(): Before returning, ensure we're at a valid position using the advance function. If row is still within bounds, there are more elements.

This design avoids extra space and efficiently skips over empty rows. It ensures each element is returned exactly once, in the correct order.

Example Walkthrough

Let's use the input vec2d = [[1,2], [], [3], [], [4,5,6]].

  • Start at row=0, col=0: vec2d[0][0]=1
  • Call next(): returns 1, advance to col=1
  • vec2d[0][1]=2: next next() returns 2, advance to col=2 (end of row), so move to row=1, col=0
  • row=1 is empty, skip to row=2, col=0
  • vec2d[2][0]=3: next() returns 3, advance to col=1 (end of row), move to row=3, col=0
  • row=3 is empty, skip to row=4, col=0
  • vec2d[4][0]=4: next() returns 4, advance to col=1
  • vec2d[4][1]=5: next() returns 5, advance to col=2
  • vec2d[4][2]=6: next() returns 6, advance to col=3 (end of row), move to row=5
  • row=5 is out of bounds, so hasNext() now returns False

The output sequence is: 1, 2, 3, 4, 5, 6. Each step shows how the iterator moves through the 2D structure, skipping empty rows as needed.

Time and Space Complexity

Brute-force approach: Flattening the entire 2D vector into a 1D list at initialization takes O(N) time and O(N) space, where N is the total number of elements.

Optimized approach (as implemented):

  • Time Complexity: Each call to next() or hasNext() may need to skip over empty rows, but over the entire lifetime of the iterator, each element and each empty row is visited at most once. Thus, the amortized time per operation is O(1).
  • Space Complexity: Only O(1) extra space is used for the row and column pointers, aside from the input.
This makes the optimized solution efficient in both time and space.

Summary

By tracking our position with row and column pointers and advancing past empty rows, we can flatten a 2D vector on-the-fly without extra storage. This approach is both elegant and efficient, requiring only O(1) extra space and providing fast iteration. The key insight is to work with the structure directly, moving forward only as needed, and skipping empty rows seamlessly.