251 Flatten 2D Vector

Implement an iterator to flatten a 2d vector.

For example, Given 2d vector =

[
  [1,2],
  [3],
  [4,5,6]
]

By callingnextrepeatedly untilhasNextreturns false, the order of elements returned bynextshould be:[1,2,3,4,5,6].

Follow up: As an added challenge, try to code it using only iterators in C++ or iterators in Java.

The Idea: Maintain a row and column pointer, and update these for each next() call. The row will increments as the column will reset to 0, when the column index is going to exceed it's current row vector. One challenge was dealing with empty vectors {}. These need to be ignored and iterated over every time we enter the start of a new row.

Complexity: O(n) time where n is the total number of elements in the vector.

class Vector2D {
public:
    Vector2D(vector<vector<int>>& vec2d) {
        v = &vec2d;
        r = c = 0;
        _skip_empty_vectors();
    }

    int next() {
        // current r and c should always be valid
        int cur_row = r, cur_col = c;

        // update for next call
        if (cur_col + 1 >= v->at(cur_row).size()) {
            c = 0;
            r++;
            _skip_empty_vectors();
        }
        else c++;

        return v->at(cur_row).at(cur_col);
    }

    bool hasNext() {
        return r < v->size();
    }

private:
    int r, c;
    vector<vector<int>> *v;

    void _skip_empty_vectors() {
        while (hasNext() && v->at(r).empty()) {
            r++;
        }
    }
};