/ Rust / 2021 / 09.rs
09.rs
 1  #![feature(test)]
 2  #![feature(binary_heap_into_iter_sorted)]
 3  
 4  use std::collections::{BinaryHeap, VecDeque};
 5  
 6  use rustc_hash::FxHashSet;
 7  
 8  type Input = Vec<Vec<u32>>;
 9  
10  fn setup(input: &str) -> Input {
11      input
12          .trim()
13          .lines()
14          .map(|line| line.trim().chars().map(|c| (c as u32) - 48).collect())
15          .collect()
16  }
17  
18  fn get_neighbors(i: i32, j: i32, w: i32, h: i32) -> impl Iterator<Item = (i32, i32)> {
19      [(-1, 0), (1, 0), (0, -1), (0, 1)]
20          .iter()
21          .cloned()
22          .filter_map(move |(dx, dy)| {
23              let x = j + dx;
24              let y = i + dy;
25              if x >= 0 && y >= 0 && x < w && y < h {
26                  Option::Some((x, y))
27              } else {
28                  Option::None
29              }
30          })
31  }
32  
33  fn part1(input: &Input) -> String {
34      let mut out = 0;
35      for (i, line) in input.iter().enumerate() {
36          for (j, c) in line.iter().enumerate() {
37              if get_neighbors(i as i32, j as i32, line.len() as i32, input.len() as i32)
38                  .all(|(p, q)| c < &input[q as usize][p as usize].clone())
39              {
40                  out += c + 1;
41              }
42          }
43      }
44      out.to_string()
45  }
46  
47  fn part2(input: &Input) -> String {
48      let mut sizes = BinaryHeap::new();
49      for (i, line) in input.iter().enumerate() {
50          for (j, c) in line.iter().enumerate() {
51              if get_neighbors(i as i32, j as i32, line.len() as i32, input.len() as i32)
52                  .all(|(p, q)| c < &input[q as usize][p as usize].clone())
53              {
54                  let mut queue = VecDeque::from([(i, j)]);
55                  let mut visited = FxHashSet::default();
56                  while !queue.is_empty() {
57                      let (y, x) = queue.pop_front().unwrap();
58  
59                      if input[y][x] == 9 {
60                          continue;
61                      }
62  
63                      if visited.contains(&(y, x)) {
64                          continue;
65                      }
66                      visited.insert((y, x));
67  
68                      for (p, q) in
69                          get_neighbors(y as i32, x as i32, line.len() as i32, input.len() as i32)
70                      {
71                          queue.push_back((q as usize, p as usize));
72                      }
73                  }
74                  sizes.push(visited.len());
75              }
76          }
77      }
78      sizes
79          .into_iter_sorted()
80          .take(3)
81          .product::<usize>()
82          .to_string()
83  }
84  
85  aoc::main!(2021, 9, ex: 1);