[LeetCode] 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 idea is very simple. We keep two variables

row

and

col

for the range of rows and cols. Specifically,

row

is the number of rows of

vec2d

and

col

is the number of columns of the current 1d vector in

vec2d

. We also keep two variables

r

and

c

to point to the current element.

In the constructor, we initialize

row

and

col

as above and initialize both

r

and

c

to be

0

(pointing to the first element).

In

hasNext()

, we just need to check whether

r

and

c

are still in the range limited by

row

and

col

.

In

next()

, we first record the current element, which is returned later. Then we update the running indexes and possibly the range if the current element is the last element of the current 1d vector.

A final and important note, since in

next()

, we record the current element, we need to guarantee that there is an element. So we implement a helper function

skipEmptyVector()

to skip the empty vectors. It is also important to handle the case that

vec2d

is empty (in this case, we set

col = -1

).

The time complexity of

hasNext()

is obviously

O(1)

and the time complexity of

next

is also

O(1)

in an amortized sense.

The code is as follows.

class Vector2D {
public:
Vector2D(vector<vector<int>>& vec2d) {
data = vec2d;
r = c = 0;
row = vec2d.size();
col = (row == 0 ? -1 : data[r].size());
skipEmptyVector();
}

int next() {
int elem = data[r][c];
if (c == col - 1) {
r++;
c = 0;
col = data[r].size();
}
else c++;
skipEmptyVector();
return elem;
}

bool hasNext() {
return col != -1 && (r < row && c < col);
}
private:
vector<vector<int>> data;
int row, col, r, c;
void skipEmptyVector(void) {
while (!col) {
r++;
col = data[r].size();
}
}
};

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

Since we need to copy the vec2d anyway, we can just copy it into a simple vector<int>, which makes life much easier :-)

class Vector2D {
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();
}
private:
vector<int> data;
int idx;
};

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

Of course, since elements in vector are stored in a contiguous range of memory, the problem can be solved in O(1) memory (see here for a better solution).