Zigzag Iterator II

LintCode Q 540 - 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”.

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]

Solution

Solution 1: Use iterators

Code:

public List<Iterator<Integer>> its;
public int turns, size, count;
    
public ZigzagIterator2(List<List<Integer>> vecs) {
	
	this.its = new ArrayList<Iterator<Integer>>();
	for (List<Integer> vec : vecs) {
		if (vec.size() > 0)
			its.add(vec.iterator());
	}
	turns = 0; count = 0;
	size = its.size();
}

public int next() {
	
	int ans = -1;  count = 0;

	if (hasNext()) {
		ans = its.get(turns % size).next();
		turns++;
	}
	
	return ans;
}

public boolean hasNext() {

	if (count == size) return false;
	
	if (its.get(turns % size).hasNext()) return true;
	
	count++; turns++;
	
	return hasNext();
}

Solution 2: A better version

Code:

public List<Iterator<Integer>> its;
public int turns;

public ZigzagIterator2(List<List<Integer>> vecs) {
	
	this.its = new ArrayList<Iterator<Integer>>();
	for (List<Integer> vec : vecs) {
		if (vec.size() > 0)
			its.add(vec.iterator());
	}
	turns = 0;
}

public int next() {
	
	int num = its.get(turns).next();

	if (its.get(turns).hasNext()) {
		turn = (turn + 1) % its.size();
	} else {
		its.remove(turns);
		if (its.size() > 0)
			turn %= its.size();
	}
}

public boolean hasNext() {
	return its.size() > 0;
}