【Leetcode】281. Zigzag Iterator

题目地址：

https://leetcode.com/problems/zigzag-iterator/

给定若干个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 ...$，也就是先尝试竖着遍历，遍历到底了再另起一列。

思路是，先按顺序存一下每一行的迭代器，然后开一个变量idx记录遍历到的list的下标。可以把主要逻辑写在hasNext里。调用hasNext的时候，先看一下idx这一行的迭代器的hasNext，如果true就直接返回true，否则向下一行走一步，也就是idx = (idx + 1) % list.size()（注意，这里取模是必要的，因为有可能下面若干行hasNext都是false，这时候相当于要回到第$0$行循环遍历），就这样一直遍历到hasNext是true的那一行，然后返回true；或者遍历回到自己的位置上，此时返回false。我们可以看到，hasNext为true的时候，idx的这一行的hasNext就是true，是有值的，所以next的逻辑就是，先调用hasNext，让idx移到下一个待遍历的位置上，再调用那一行的迭代器的next。代码如下（题目里初始化参数是两个列表，但上面的算法适合于多个列表的情形）：

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

public class ZigzagIterator {
    
    private List<Iterator<Integer>> list;
    // 标记当前遍历到的list的下标
    private int idx;
    
    public ZigzagIterator(List<Integer> v1, List<Integer> v2) {
    	// 将两个列表的迭代器加入到list里去
        list = new ArrayList<>();
        list.add(v1.iterator());
        list.add(v2.iterator());
    }
    
    public int next() {
        hasNext();
        int next = list.get(idx).next();
        idx = (idx + 1) % list.size();
        return next;
    }
    
    public boolean hasNext() {
    	// 如果idx这一行有值，那就返回true
        if (list.get(idx).hasNext()){
            return true;
        }
        
        // 否则缓存一下idx赋值给tmp，然后让idx后移，
        // 直到发现了hasNext的那一行返回true，或者回到tmp返回false
        int tmp = idx;
        // 先向下走一步，不然进不了循环。当然用do-while循环也可以
        idx = (idx + 1) % list.size();
        while (idx != tmp) {
            if (list.get(idx % list.size()).hasNext()) {
                return true;
            }
            
            // 向下走一步的时候需要取模，也就是要循环遍历
            idx = (idx + 1) % list.size();
        }
        
        return false;
    }
}

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