[LintCode] Problem 601 - Flatten 2D Vector

Posted on | In ,

Implement an iterator to flatten a 2d vector.

Example

No.1

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

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

No.2

Input:[[7,9],[5]]

Output:[7,9,5]

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
private Iterator<List<Integer>> row;
private Iterator<Integer> col;

public Vector2D(List<List<Integer>> vec2d) {
    row = vec2d.iterator();
}

@Override
public Integer next() {
    if (!hasNext())
        return null;

    return col.next();
}

@Override
public boolean hasNext() {
    while ((col == null || !col.hasNext()) && row.hasNext())
        col = row.next().iterator();

    return col != null && col.hasNext();
}