rustgym/leetcode/
_407_trapping_rain_water_2.rs

1struct Solution;
2
3use std::cmp::Reverse;
4use std::collections::BinaryHeap;
5
6impl Solution {
7    fn trap_rain_water(height_map: Vec<Vec<i32>>) -> i32 {
8        let n = height_map.len();
9        let m = height_map[0].len();
10        let mut queue: BinaryHeap<(Reverse<i32>, usize, usize)> = BinaryHeap::new();
11        let mut visited = vec![vec![false; m]; n];
12        for i in 0..n {
13            for j in 0..m {
14                if i == 0 || j == 0 || i == n - 1 || j == m - 1 {
15                    visited[i][j] = true;
16                    queue.push((Reverse(height_map[i][j]), i, j));
17                }
18            }
19        }
20        let mut res = 0;
21        while let Some((Reverse(h), i, j)) = queue.pop() {
22            if i > 0 && !visited[i - 1][j] {
23                let ii = i - 1;
24                visited[ii][j] = true;
25                if h > height_map[ii][j] {
26                    res += h - height_map[ii][j];
27                    queue.push((Reverse(h), ii, j));
28                } else {
29                    queue.push((Reverse(height_map[ii][j]), ii, j));
30                }
31            }
32            if j > 0 && !visited[i][j - 1] {
33                let jj = j - 1;
34                visited[i][jj] = true;
35                if h > height_map[i][jj] {
36                    res += h - height_map[i][jj];
37                    queue.push((Reverse(h), i, jj));
38                } else {
39                    queue.push((Reverse(height_map[i][jj]), i, jj));
40                }
41            }
42            if i + 1 < n && !visited[i + 1][j] {
43                let ii = i + 1;
44                visited[ii][j] = true;
45                if h > height_map[ii][j] {
46                    res += h - height_map[ii][j];
47                    queue.push((Reverse(h), ii, j));
48                } else {
49                    queue.push((Reverse(height_map[ii][j]), ii, j));
50                }
51            }
52            if j + 1 < m && !visited[i][j + 1] {
53                let jj = j + 1;
54                visited[i][jj] = true;
55                if h > height_map[i][jj] {
56                    res += h - height_map[i][jj];
57                    queue.push((Reverse(h), i, jj));
58                } else {
59                    queue.push((Reverse(height_map[i][jj]), i, jj));
60                }
61            }
62        }
63        res
64    }
65}
66
67#[test]
68fn test() {
69    let height_map = vec_vec_i32![[1, 4, 3, 1, 3, 2], [3, 2, 1, 3, 2, 4], [2, 3, 3, 2, 3, 1]];
70    let res = 4;
71    assert_eq!(Solution::trap_rain_water(height_map), res);
72}
rustgym/leetcode/_407_trapping_rain_water_2.rs

rustgym/leetcode/
_407_trapping_rain_water_2.rs
