rustgym/leetcode/
_305_number_of_islands_2.rs

1struct Solution;
2
3struct UnionFind {
4    parent: Vec<usize>,
5    n: usize,
6}
7
8impl UnionFind {
9    fn new(n: usize) -> Self {
10        let parent = (0..n).collect();
11        UnionFind { parent, n }
12    }
13    fn find(&mut self, i: usize) -> usize {
14        let j = self.parent[i];
15        if i == j {
16            i
17        } else {
18            let k = self.find(j);
19            self.parent[i] = k;
20            k
21        }
22    }
23
24    fn union(&mut self, i: usize, j: usize) -> bool {
25        let i = self.find(i);
26        let j = self.find(j);
27        if i != j {
28            self.parent[i] = j;
29            self.n -= 1;
30            true
31        } else {
32            false
33        }
34    }
35}
36
37impl Solution {
38    fn num_islands2(m: i32, n: i32, positions: Vec<Vec<i32>>) -> Vec<i32> {
39        let m = m as usize;
40        let n = n as usize;
41        let mut uf = UnionFind::new(m * n);
42        let mut grid = vec![vec![0; n]; m];
43        let mut group = 0;
44        let mut res = vec![];
45        for position in positions {
46            let r = position[0] as usize;
47            let c = position[1] as usize;
48            let i = r * n + c;
49            if grid[r][c] == 1 {
50                res.push(group);
51                continue;
52            }
53            grid[r][c] = 1;
54            group += 1;
55            if r > 0 && grid[r - 1][c] == 1 {
56                let j = (r - 1) * n + c;
57                if uf.union(i, j) {
58                    group -= 1;
59                }
60            }
61            if c > 0 && grid[r][c - 1] == 1 {
62                let j = r * n + c - 1;
63                if uf.union(i, j) {
64                    group -= 1;
65                }
66            }
67            if r + 1 < m && grid[r + 1][c] == 1 {
68                let j = (r + 1) * n + c;
69                if uf.union(i, j) {
70                    group -= 1;
71                }
72            }
73            if c + 1 < n && grid[r][c + 1] == 1 {
74                let j = r * n + c + 1;
75                if uf.union(i, j) {
76                    group -= 1;
77                }
78            }
79            res.push(group);
80        }
81        res
82    }
83}
84
85#[test]
86fn test() {
87    let m = 3;
88    let n = 3;
89    let positions = vec_vec_i32![[0, 0], [0, 1], [1, 2], [2, 1]];
90    let res = vec![1, 1, 2, 3];
91    assert_eq!(Solution::num_islands2(m, n, positions), res);
92}