LintCode 541: Zigzag Iterator II

541. Zigzag Iterator II

Follow up Zigzag Iterator: What if you are given k 1d vectors? How well can your code be extended to such cases? The “Zigzag” order is not clearly defined and is ambiguous for k > 2 cases. If “Zigzag” does not look right to you, replace “Zigzag” with “Cyclic”.

Example
Example1

Input: k = 3
vecs = [
[1,2,3],
[4,5,6,7],
[8,9],
]
Output: [1,4,8,2,5,9,3,6,7]
Example2

Input: k = 3
vecs = [
[1,1,1]
[2,2,2]
[3,3,3]
]
Output: [1,2,3,1,2,3,1,2,3]

解法1：用轮询方式。
代码如下：

class ZigzagIterator2 {
public:
    /*
    * @param vecs: a list of 1d vectors
    */
    ZigzagIterator2(vector<vector<int>>& vecs) {
        vecsSize = vecs.size();
        pollIndex = -1;
        iters.resize(vecsSize);
        iterEnds.resize(vecsSize);
        for (int i = 0; i < vecsSize; ++i) {
            iters[i] = vecs[i].begin();
            iterEnds[i] = vecs[i].end();
        }
    }

    /*
     * @return: An integer
     */
    int next() {
        if (hasNext()) {

            for (int i = pollIndex + 1; i < vecsSize; ++i) {
                if (iters[i] != iterEnds[i]) {
                    pollIndex = i;
                    return *(iters[i]++);
                }
            }
            
            for (int i = 0; i <= pollIndex; ++i) {
                if (iters[i] != iterEnds[i]) {
                    pollIndex = i;            
                    return *(iters[i]++);
                }
            }
        }
    }

    /*
     * @return: True if has next
     */
    bool hasNext() {
        for (int i = max(0, pollIndex); i < vecsSize; ++i) {
            if (iters[i] != iterEnds[i]) return true;
        }
        for (int i = 0; i <= pollIndex; ++i) {
            if (iters[i] != iterEnds[i]) return true;
        }
        return false;
    }
    
private:
    vector<vector<int>::iterator> iters, iterEnds;
    
    //vector<vector<int>>::iterator pollIter, pollIterEnd;
    int pollIndex, vecsSize;
};
/**
 * Your ZigzagIterator2 object will be instantiated and called as such:
 * ZigzagIterator2 solution(vecs);
 * while (solution.hasNext()) result.push_back(solution.next());
 * Ouptut result
 */