[LeetCode] Zigzag Iterator

Zigzag Iterator

Given two 1d vectors, implement an iterator to return their elements alternately.

For example, given two 1d vectors:

v1 = [1, 2]
v2 = [3, 4, 5, 6]

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

Follow up: What if you are given k 1d vectors? How well can your code be extended to such cases?

后天Onsite，求人品，求Offer！ ( T,T

class ZigzagIterator {
private:
    vector<int> &mv1, &mv2;
    vector<int>::iterator it1, it2, it;
public:
    ZigzagIterator(vector<int>& v1, vector<int>& v2) : mv1(v1), mv2(v2) {
        it1 = v1.begin();
        it2 = v2.begin();
        it = (it1 == mv1.end()) ? it2 : it1;
    }

    int next() {
        int val = *it;
        if (it == it1) {
            ++it1;
            it = (it2 == mv2.end()) ? it1 : it2;
        } else {
            ++it2;
            it = (it1 == mv1.end()) ? it2 : it1;
        }
        return val;
    }

    bool hasNext() {
        return it1 != mv1.end() || it2 != mv2.end();
    }
};

/**
 * Your ZigzagIterator object will be instantiated and called as such:
 * ZigzagIterator i(v1, v2);
 * while (i.hasNext()) cout << i.next();
 */