Leetcode-231.Flatten 2D Vector

Problem Description:

　　Implement an iterator to flatten a 2d vector.

For example,
Given 2d vector =

[
[1,2],
[3],
[4,5,6]
]
By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1, 2, 3, 4, 5, 6].

The initial code :

class Vector2D {
public:
    Vector2D(vector<vector<int>>& vec2d) {

    }

    int next() {

    }

    bool hasNext() {

    }
};
/**
 * Your Vector2D object will be instantiated and called as such:
 * Vector2D i(vec2d);
 * while (i.hasNext()) cout << i.next();
 *
*/

---

Analysis:
We could simply save the elements to a 1D vector in the constructor.

class Vector2D {
private:
    vector<int> data;
    int idx;
public:
    Vector2D(vector<vector<int>>& vec2d) {
      int row = vec2d.size();
      for (int r = 0; r < row; ++r)
      {
        int col = vec2d[r].size();
        for (int c = 0; c < col; c++)
                data.push_back(vec2d[r][c]);

      }
      idx = 0;
    }

    int next() {
        return data[idx++];
    }

    bool hasNext() {
        return idx < data.size();
    }
};

Short code in inserting the element to vector.

 for (auto& v : vec2d)
            data.insert(end(data), begin(v), end(v));

---

O(1) space :

class Vector2D {
private :
    vector<vector<int> > :: iterator i, iEnd;
    int j = 0;
public:
    Vector2D(vector<vector<int>>& vec2d) {
        i = vec2d.begin();
        iEnd = vec2d.end();
    }

    int next() {
        hasNext();
        return (*i)[j++]
    }

    bool hasNext() {
        while (i != iEnd && j == (*i).size())
            i++, j = 0;
        return i != iEnd;
    }
};