Grid Illumination LeetCode solution in Rust

1
use std::collections::hash_map::{Entry, OccupiedEntry};
2
use std::collections::{HashMap, HashSet};
3

4
impl Solution {
5
  pub fn grid_illumination(n: i32, lamps: Vec<Vec<i32>>, queries: Vec<Vec<i32>>) -> Vec<i32> {
6
    let (mut rows, mut cols, mut digL, mut digR, mut lamps) = lamps.iter().fold(
7
      (
8
        HashMap::new(),
9
        HashMap::new(),
10
        HashMap::new(),
11
        HashMap::new(),
12
        HashSet::new(),
13
      ),
14
      |(mut rows, mut cols, mut digL, mut digR, mut lamps), l| {
15
        if lamps.insert((l[0], l[1])) {
16
          *rows.entry(l[0]).or_default() += 1;
17
          *cols.entry(l[1]).or_default() += 1;
18
          *digL.entry(l[0] - l[1]).or_default() += 1;
19
          *digR.entry(l[0] + l[1]).or_default() += 1;
20
        }
21
        (rows, cols, digL, digR, lamps)
22
      },
23
    );
24

25
    queries
26
      .iter()
27
      .map(|q| {
28
        if !rows.contains_key(&q[0])
29
          && !cols.contains_key(&q[1])
30
          && !digL.contains_key(&(q[0] - q[1]))
31
          && !digR.contains_key(&(q[0] + q[1]))
32
        {
33
          return 0;
34
        }
35
        for i in (q[0] - 1).max(0)..=(q[0] + 1).min(n - 1) {
36
          for j in ((q[1] - 1).max(0)..=(q[1] + 1).min(n - 1)).filter(|&j| lamps.remove(&(i, j))) {
37
            match (
38
              rows.entry(i),
39
              cols.entry(j),
40
              digL.entry(i - j),
41
              digR.entry(i + j),
42
            ) {
43
              (
44
                Entry::Occupied(ro),
45
                Entry::Occupied(co),
46
                Entry::Occupied(dl),
47
                Entry::Occupied(dr),
48
              ) => {
49
                let arr = [ro, co, dl, dr];
50
                for mut e in arr {
51
                  let mut v = e.get_mut();
52
                  match *v {
53
                    1 => {
54
                      e.remove();
55
                    }
56
                    _ => *v -= 1,
57
                  }
58
                }
59
              }
60
              _ => unreachable!(),
61
            }
62
          }
63
        }
64
        1
65
      })
66
      .collect()
67
  }
68
}