Problem description:

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

Example:

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

Output: [1,3,2,4,5,6]

Explanation: 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?

Clarification for the follow up question:
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”. For example:

1
2
3
4
5
6
Input:
[1,2,3]
[4,5,6,7]
[8,9]

Output: [1,4,8,2,5,9,3,6,7].

Solution:

The idea of this problem is to use a queue to store the current iterator of each vector. At start, we need to store the begin and end iterator of each vector. Every time we need to output a value, we take the pair<begin, end> from the queue. After we get the pair, we can use the begin iterator to find next element in that vector.

  1. Use a queue to store iterators of each vector, push pairs of iterators with the zigzag order.
  2. Take the front element of the queue, pop it out. If we add 1 to the begin iterator, then we can get to the next element in the vector. However, if the begin iterator is equal to the end iterator, that means we reach to the end. And we don’t need to push the pair in to queue again.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class ZigzagIterator {
public:
    ZigzagIterator(vector<int>& v1, vector<int>& v2) {
        if(!v1.empty()) q.push(make_pair(v1.begin(), v1.end()));
        if(!v2.empty()) q.push(make_pair(v2.begin(), v2.end()));
        //if there're multiple vectors, then push it into the queue
    }

    int next() {
        auto it= q.front().first, end= q.front().second;
        q.pop();
        if(it+1 != end) q.push(make_pair(it+1, end));
        return *it;
    }

    bool hasNext() {
        return !q.empty();
    }
    
    queue<pair<vector<int>::iterator, vector<int>::iterator>> q;
};

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

time complexity: $O(n)$
space complexity: $O(n)$
reference: