LeetCode - Zigzag Iterator

ZigzagIterator.java

// 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?

// Clarification for the follow up question - Update (2015-09-18):
// 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, given the following input:

// [1,2,3]
// [4,5,6,7]
// [8,9]
// It should return [1,4,8,2,5,9,3,6,7].

public class ZigzagIterator {
    
    private List<Integer> res;
    private int index = 0;

    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
        res = new ArrayList<Integer>();
        int len1 = v1.size(), len2 = v2.size();
        int i = 0;
        while(i < len1 && i < len2){
            res.add(v1.get(i));
            res.add(v2.get(i));
            i++;
        }
        
        while(i < len1) res.add(v1.get(i++));
        while(i < len2) res.add(v2.get(i++));
    }

    public int next() {
        return res.get(index++);
    }

    public boolean hasNext() {
        return index != res.size();
    }
}

/**
 * Your ZigzagIterator object will be instantiated and called as such:
 * ZigzagIterator i = new ZigzagIterator(v1, v2);
 * while (i.hasNext()) v[f()] = i.next();
 */
 
// solution 2 --- use queue
 
// solution 3 --- O(1) space https://leetcode.com/discuss/58012/short-java-o-1-space
public class ZigzagIterator {
    Iterator<Integer> it1 = null, it2 = null;
    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
        if(v1 != null) it1 = v1.iterator();
        if(v2 != null) it2 = v2.iterator();
    }

    public int next() {
        int val = -1;
        if(hasNext()){
            if(it1.hasNext()){
                val = it1.next();
                swap();
            } else{
                swap();
                val = it1.next();
            }
        }
        return val;
    }
    
    private void swap(){
        Iterator<Integer> temp = it1;
        it1 = it2;
        it2 = temp;
    }

    public boolean hasNext() {
        return it1.hasNext() || it2.hasNext();
    }
}