[LeetCode] Zigzag Iterator

Problem Description:

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?

Update: when k vectors are provided, what the results should be still remain to be a question (you may refer to this post). So the following notes do not focus on that now.

The idea is as follows: keep the two beginning iterators and the two end iterators of v1 and v2into two arrays of type vector<int>::iterator with size 2. Then we keep a variable p(initialized to be 0) to record which iterator should be used: p takes values ranging from 0 to1. Each time we call next, update p by p = (p + 1) % 2 to circulate it to point to the next vector.

The code is as follows.

 1 class ZigzagIterator {
 2 public:
 3     ZigzagIterator(vector<int>& v1, vector<int>& v2) {
 4         bs[0] = v1.begin(), bs[1] = v2.begin();
 5         es[0] = v1.end(), es[1] = v2.end();
 6         p = 0;
 7     }
 8 
 9     int next() {
10         int elem;
11         if (bs[0] == es[0]) elem = *bs[1]++;
12         else if (bs[1] == es[1]) elem = *bs[0]++;
13         else {
14             elem = *bs[p]++;
15             p = (p + 1) % 2;
16         }
17         return elem;
18     }
19 
20     bool hasNext() {
21         return bs[0] != es[0] || bs[1] != es[1];
22     }
23 private:
24     int p;
25     vector<int>::iterator bs[2], es[2];
26 };