Flatten 2D Vector

Implement an iterator to flatten a 2d vector.

For example, Given 2d vector =

[
  [1,2],
  [3],
  [4,5,6]
] 

By calling next repeatedly until hasNext returns false, the order of elements returned by next should be: [1,2,3,4,5,6].

数组法

复杂度

时间 O(N) 空间 O(1)

思路

用一个数组表示每个List的迭代器，然后再记录一个变量，用来表示当前用到了第几个迭代器。

代码

public class Vector2D {

    List<Iterator<Integer>> its;
    int curr = 0;
    
    public Vector2D(List<List<Integer>> vec2d) {
        this.its = new ArrayList<Iterator<Integer>>();
        for(List<Integer> l : vec2d){
            // 只将非空的迭代器加入数组
            if(l.size() > 0){
               this.its.add(l.iterator()); 
            }
        }
    }

    public int next() {
        Integer res = its.get(curr).next();
        // 如果该迭代器用完了，换到下一个
        if(!its.get(curr).hasNext()){
            curr++;
        }
        return res;
    }

    public boolean hasNext() {
        return curr < its.size() && its.get(curr).hasNext();
    }
}

双迭代器法

复杂度

时间 O(N) 空间 O(1)

思路

维护两个迭代器：一个是输入的List<List<Integer>>的迭代器，它负责遍历List<Integer>的迭代器。另一个则是List<Integer>的迭代器，它负责记录当前到哪一个List的迭代器了。每次next时，我们先调用一下hasNext，确保当前List的迭代器有下一个值。

代码

public class Vector2D {

    Iterator<List<Integer>> it;
    Iterator<Integer> curr;
    
    public Vector2D(List<List<Integer>> vec2d) {
        it = vec2d.iterator();
    }

    public int next() {
        hasNext();
        return curr.next();
    }

    public boolean hasNext() {
        // 当前列表的迭代器为空，或者当前迭代器中没有下一个值时，需要更新为下一个迭代器
        while((curr == null || !curr.hasNext()) && it.hasNext()){
            curr = it.next().iterator();
        }
        return curr != null && curr.hasNext();
    }
}