Skip to main content

rustgym/leetcode/
_251_flatten_2d_vector.rs

1struct Vector2D {
2    array: Vec<i32>,
3    count: usize,
4}
5
6impl Vector2D {
7    fn new(v: Vec<Vec<i32>>) -> Self {
8        let mut array = vec![];
9        for row in v {
10            for col in row {
11                array.push(col);
12            }
13        }
14        Vector2D { array, count: 0 }
15    }
16    fn next(&mut self) -> i32 {
17        let res = self.array[self.count];
18        self.count += 1;
19        res
20    }
21    fn has_next(&self) -> bool {
22        self.count < self.array.len()
23    }
24}
25
26#[test]
27fn test() {
28    let v = vec![vec![1, 2], vec![3], vec![4]];
29    let mut obj = Vector2D::new(v);
30    assert_eq!(obj.next(), 1);
31    assert_eq!(obj.next(), 2);
32    assert_eq!(obj.next(), 3);
33    assert_eq!(obj.has_next(), true);
34    assert_eq!(obj.has_next(), true);
35    assert_eq!(obj.next(), 4);
36    assert_eq!(obj.has_next(), false);
37}