Flatten 2D Vector

https://www.lintcode.com/problem/flatten-2d-vector/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].

Follow up:

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

Thoughts

把二维数组看作一维，实现iterator。i和j表当前行和列，调用hasNext时检查j此时是否超出data[i]的size，是就循环直到把i和j移到下一个合法元素；如果当i超出M时说明没有下一个元素。调用next()时先调用hasNext移动指针，再返回结果和++j。

Code

class Vector2D {
public:
    vector<vector<int>> *ptr_vec2d;
    int i = 0, j = 0;
    Vector2D(vector<vector<int>>& vec2d) {
        // Initialize your data structure here
        ptr_vec2d = &vec2d;
        hasNext();
    }

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

    bool hasNext() {
        while (i < ptr_vec2d->size() && j == (*ptr_vec2d)[i].size()) {
            ++i;
            j = 0;
        }
        return i < ptr_vec2d->size();
    }
};

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

public class Vector2D implements Iterator<Integer> {
    Iterator<Integer> inIter;
    Iterator<List<Integer>> outIter;

    public Vector2D(List<List<Integer>> vec2d) {
        outIter = vec2d.iterator();
    }

    @Override
    public Integer next() {
        hasNext();
        return inIter.next();
    }

    @Override
    public boolean hasNext() {
        while (outIter.hasNext() && (inIter == null || !inIter.hasNext())) {
            inIter = outIter.next().iterator();
        } 
        return inIter == null ? false : inIter.hasNext();
    }
}

/**
 * Your Vector2D object will be instantiated and called as such:
 * Vector2D i = new Vector2D(vec2d);
 * while (i.hasNext()) v[f()] = i.next();
 */

Analysis

时间复杂度O(N).