【Lintcode】541. Zigzag Iterator II

题目地址：

https://www.lintcode.com/problem/zigzag-iterator-ii/description

给定若干个list，要求构造一个迭代器，它的遍历顺序是一列一列遍历。这里的一列一列遍历的意思是，比如说将这些list作为行排成二维数组的话（注意，这里每一行未必长度相等），那么可以写成：$$A=\left[\begin{array}{lccc}{a}_{00}& a_{01}& ...& a_{0s_0}\\ a_{10}& a_{11}& ...& a_{1s_1}\\ ...& & & \\ a_{k0}& a_{k1}& ...& a_{k{s}_k}\end{array}\right]$$那么遍历顺序应该是$a_{00}\to a_{10}\to ...\to a_{k0}\to a_{01}\to a_{11}\to ...$，也就是先尝试竖着遍历，遍历到底了再另起一列。

思路是，先把每个list的迭代器存进一个list里，另外再开个变量保存当前遍历到了哪个list了。hasNext其实就是在找哪一行的迭代器的hasNext为true。调用next的时候，只需要先调用hasNext，将光标移到hasNext为true的那一行，接着调用next即可，注意，此时需要将光标向下移动一行，否则下次调用next的时候就是遍历的之前刚遍历的那一行了。代码如下：

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

public class ZigzagIterator2 {
    
    private List<Iterator<Integer>> list;
    private int idx;
    
    /*
     * @param vecs: a list of 1d vectors
     */
    public ZigzagIterator2(List<List<Integer>> vecs) {
        // do intialization if necessary
        list = new ArrayList<>();
        for (List<Integer> vec : vecs) {
            list.add(vec.iterator());
        }
    }
    
    /*
     * @return: An integer
     */
    public int next() {
        // write your code here
        hasNext();
        // 把结果缓存一下
        int res = list.get(idx).next();
        // 将idx下移一格
        idx = (idx + 1) % list.size();
        return res;
    }
    
    /*
     * @return: True if has next
     */
    public boolean hasNext() {
        // write your code here
        // 如果当前行还有值，则返回true
        if (list.get(idx).hasNext()) {
            return true;
        }
        
        int tmp = idx;
        // 否则向后循环遍历，直到遍历到有值的一行，或者回到起点
        idx = (idx + 1) % list.size();
        while (idx != tmp) {
            if (list.get(idx).hasNext()) {
                return true;
            }
            idx = (idx + 1) % list.size();
        }
        
        return false;
    }
}

两个操作时间复杂度是$O(n)$，$n$为行数，空间也是$O(n)$。这个时间复杂度看起来很高，但是由于一旦调用了hasNext之后，迭代器就不会回头了，所以每个数字其实只会遍历一次，真实耗费时间并没有那么多。