Add day09, 2021

3 files changed, 244 insertions, 1 deletions

diff --git a/2021/rust/Cargo.toml b/2021/rust/Cargo.toml
index 3f73cb5..dd0516c 100644
--- a/2021/rust/Cargo.toml
+++ b/2021/rust/Cargo.toml
@@ -8,7 +8,7 @@ members = [
     "day06",
     "day07",
     "day08",
-#    "day09",
+    "day09",
 #    "day10",
 #    "day11",
 #    "day12",
diff --git a/2021/rust/day09/Cargo.toml b/2021/rust/day09/Cargo.toml
new file mode 100644
index 0000000..e8bb2c8
--- /dev/null
+++ b/2021/rust/day09/Cargo.toml
@@ -0,0 +1,9 @@
+[package]
+name = "day09"
+version = "0.1.0"
+authors = ["cos <cos>"]
+edition = "2021"
+
+[dependencies]
+aoc = { path = "../../../common/rust/aoc" }
+anyhow = "1.0"
diff --git a/2021/rust/day09/src/main.rs b/2021/rust/day09/src/main.rs
new file mode 100644
index 0000000..8bf14c1
--- /dev/null
+++ b/2021/rust/day09/src/main.rs
@@ -0,0 +1,234 @@
+use {
+    anyhow::{
+        anyhow,
+        Context,
+        Result,
+    },
+    std::{
+        env::args,
+        fs::File,
+        io::{
+            BufRead,
+            BufReader,
+        },
+        path::Path,
+    },
+};
+
+#[derive(Clone,Copy,Debug)]
+enum Direction {
+    North,
+    South,
+    East,
+    West,
+}
+
+#[derive(Clone,Copy,Debug)]
+struct Point {
+    x: usize,
+    y: usize,
+}
+
+impl Point {
+    fn neighbour(&self, dir: Direction) -> Point {
+        match dir {
+            Direction::North => Point {x: self.x, y: self.y.overflowing_sub(1).0},
+            Direction::South => Point {x: self.x, y: self.y + 1},
+            Direction::East  => Point {x: self.x + 1, y: self.y},
+            Direction::West  => Point {x: self.x.overflowing_sub(1).0, y: self.y},
+        }
+    }
+}
+
+#[derive(Debug)]
+struct HeightMap {
+    size: Point,
+    levels: Vec<u8>,
+}
+
+#[derive(Debug)]
+struct VisitMap {
+    size: Point,
+    visited: Vec<bool>,
+}
+
+impl VisitMap {
+    fn new(size: Point) -> Self{
+        let visited = vec![false; size.x*size.y];
+        Self {
+            size,
+            visited,
+        }
+    }
+
+    fn is_visited(&self, pos: Point) -> bool {
+        if pos.x >= self.size.x || pos.y >= self.size.y {
+            return true;
+        }
+
+        self.visited.get(pos.y * self.size.x + pos.x).copied().expect("Invalid index")
+    }
+
+    fn visit(&mut self, pos: Point) {
+        self.visited[pos.y * self.size.x + pos.x] = true;
+    }
+}
+
+impl HeightMap {
+    fn level(&self, pos: &Point) -> Result<u8> {
+        if pos.x >= self.size.x || pos.y >= self.size.y {
+            return Err(anyhow!("Invalid index {}, {}", pos.x, pos.y));
+        }
+        self.levels.get(pos.y * self.size.x + pos.x).copied()
+            .ok_or_else(|| anyhow!("Invalid index {}, {}", pos.x, pos.y))
+    }
+
+    fn is_lowpoint(&self, pos: &Point) -> bool {
+        let us = self.level(pos).unwrap_or(0);
+
+        let mut ok_count = 0;
+        let directions = [Direction::North, Direction::South, Direction::East, Direction::West];
+
+        directions.into_iter().for_each(|dir| {
+            let neighbour = pos.neighbour(dir);
+            match self.level(&neighbour) {
+                Ok(them) => if us < them { ok_count += 1; },
+                Err(_) => { ok_count += 1; },
+            }
+        });
+        ok_count == directions.into_iter().count()
+    }
+
+    fn is_highpoint(&self, pos: &Point) -> bool {
+        self.level(pos).unwrap_or(9) == 9
+    }
+
+    fn basin_size(&self, start: &Point) -> Result<usize> {
+        let mut unexplored = vec![*start];
+        let mut pos = None;
+        let mut size = 0;
+
+        let mut seen = VisitMap::new(self.size);
+
+        while !unexplored.is_empty() {
+            if pos.is_none() {
+                if let Some(new_pos) = unexplored.pop() {
+                    pos = Some(new_pos);
+                    seen.visit(new_pos);
+                    size += 1;
+                } else {
+                    unreachable!();
+                }
+            }
+
+            if let Some(cur_pos) = pos {
+                for dir in [Direction::North, Direction::South, Direction::East, Direction::West] {
+                    if seen.is_visited(cur_pos.neighbour(dir)) {
+                        continue;
+                    }
+
+                    let neighbour = cur_pos.neighbour(dir);
+                    if self.is_highpoint(&neighbour) {
+                        seen.visit(neighbour);
+                        continue;
+                    } else {
+                        seen.visit(neighbour);
+                        unexplored.push(neighbour);
+                        continue;
+                    }
+                }
+                if seen.is_visited(cur_pos.neighbour(Direction::North)) &&
+                    seen.is_visited(cur_pos.neighbour(Direction::South)) &&
+                    seen.is_visited(cur_pos.neighbour(Direction::East)) &&
+                    seen.is_visited(cur_pos.neighbour(Direction::West))
+                {
+                    pos = None;
+                }
+            } else {
+                unreachable!();
+            }
+        }
+
+        Ok(size)
+    }
+
+    fn risk_level(&self, pos: &Point) -> Result<usize> {
+        let level = self.level(pos)?;
+
+        Ok(level as usize + 1)
+    }
+}
+
+fn read_input<T: AsRef<Path>>(filename: T) -> Result<HeightMap> {
+    let reader = BufReader::new(File::open(filename)?);
+    let mut x_opt = None;
+    let mut y = 0;
+
+    let levels = reader.lines().map(|input| {
+        let i = input?;
+        let o: Vec<u8> = i.chars().map(|column| {
+            let level: u8 = column.to_digit(10).ok_or(anyhow!("Invalid digit"))? as u8;
+            Ok(level)
+        }).collect::<Result<Vec<u8>>>()?;
+        match (x_opt, o.len()) {
+            (None, len) => x_opt = Some(len),
+            (Some(this), other) => if this != other {
+                return Err(anyhow!("Inconsistent map width encountered: {} != {}", this, other));
+            },
+        }
+        y += 1;
+        Ok(o)
+    }).collect::<Result<Vec<Vec<u8>>>>()?.into_iter().flatten().collect();
+
+    let x = x_opt.ok_or(anyhow!("Could not determine size of heightmap."))?;
+    let size = Point {
+        x,
+        y,
+    };
+    Ok(HeightMap {
+        size,
+        levels,
+    })
+}
+
+fn part1(map: &HeightMap) -> Result<usize> {
+    let mut risk_level = 0;
+    for y in 0..map.size.y {
+        for x in 0..map.size.x {
+            if map.is_lowpoint(&Point {x, y}) {
+                risk_level += map.risk_level(&Point {x, y})?;
+            }
+        }
+    }
+    Ok(risk_level)
+}
+
+fn part2(map: &HeightMap) -> Result<usize> {
+    let mut basin_sizes = vec![];
+    for y in 0..map.size.y {
+        for x in 0..map.size.x {
+            if map.is_lowpoint(&Point {x, y}) {
+                basin_sizes.push(map.basin_size(&Point {x, y})?);
+            }
+        }
+    }
+    basin_sizes.sort_unstable();
+
+    Ok(basin_sizes.iter().rev().take(3).product())
+}
+
+fn main() -> Result<()> {
+    let ( do_part_1, do_part_2 ) = aoc::do_parts();
+
+    let filename = args().nth(1).ok_or(anyhow!("Missing input filename"))?;
+    let input = read_input(filename).context("Could not read input")?;
+    if do_part_1 {
+        let solution = part1(&input).context("No solution for part 1")?;
+        println!("Part1, solution found to be: {}", solution);
+    }
+    if do_part_2 {
+        let solution = part2(&input).context("No solution for part 2")?;
+        println!("Part2, solution found to be: {}", solution);
+    }
+    Ok(())
+}