Implement 2023 day 22 part 2

2023/src/day22.rs

@@ -1,7 +1,166 @@
-pub fn part1(_input: &[u8]) -> anyhow::Result<String> {
-    anyhow::bail!("Not implemented")
+use std::collections::HashSet;
+
+use nom::bytes::complete::tag;
+use nom::combinator::map;
+use nom::multi::many1;
+use nom::sequence::terminated;
+use nom::sequence::tuple;
+use nom::IResult;
+
+use crate::common::minmax;
+use crate::common::parse_input;
+
+struct Brick([[u16; 3]; 2]);
+
+impl Brick {
+    fn lowest(&self) -> u16 {
+        Ord::min(self.0[0][2], self.0[1][2])
+    }
+
+    fn iter_squares(&self) -> impl Iterator<Item = (u16, u16)> {
+        let xs = minmax(self.0[0][0], self.0[1][0]);
+        let ys = minmax(self.0[0][1], self.0[1][1]);
+
+        (ys.0..=ys.1).flat_map(move |z| (xs.0..=xs.1).map(move |x| (x, z)))
+    }
+
+    fn parse(i: &[u8]) -> IResult<&[u8], Brick> {
+        use nom::character::complete::u16;
+        let parse_coordinates = |i| {
+            map(
+                tuple((terminated(u16, tag(",")), terminated(u16, tag(",")), u16)),
+                |(x, y, z)| [x, y, z],
+            )(i)
+        };
+
+        map(
+            tuple((
+                terminated(parse_coordinates, tag("~")),
+                terminated(parse_coordinates, tag("\n")),
+            )),
+            |(first, second)| Brick([first, second]),
+        )(i)
+    }
 }
 
-pub fn part2(_input: &[u8]) -> anyhow::Result<String> {
-    anyhow::bail!("Not implemented")
+fn parse_bricks(i: &[u8]) -> IResult<&[u8], Vec<Brick>> {
+    many1(Brick::parse)(i)
+}
+
+fn compute_support(input: &[u8]) -> anyhow::Result<(Vec<Vec<usize>>, Vec<u8>)> {
+    let mut bricks = parse_input(input, parse_bricks)?;
+    bricks.sort_unstable_by_key(Brick::lowest);
+
+    let mut width = 0;
+    let mut breadth = 0;
+
+    for brick in &bricks {
+        width = Ord::max(width, Ord::max(brick.0[0][0], brick.0[1][0]));
+        breadth = Ord::max(breadth, Ord::max(brick.0[0][2], brick.0[1][2]));
+    }
+
+    let width = width as usize + 1;
+    let breadth = breadth as usize + 1;
+
+    let mut height_map = vec![0; width * breadth];
+    let mut top_map = vec![usize::MAX; height_map.len()];
+    let mut supported = vec![0u8; bricks.len()];
+    let mut supporting = vec![Vec::new(); bricks.len()];
+
+    let mut temp = HashSet::new();
+
+    for ((i, brick), supported) in bricks.iter().enumerate().zip(&mut supported) {
+        let max_z = brick
+            .iter_squares()
+            .map(|(x, y)| height_map[x as usize + (y as usize) * width])
+            .max()
+            .expect("Guaranteed non-empty iterator.");
+
+        let zs = minmax(brick.0[0][2], brick.0[1][2]);
+
+        debug_assert!(
+            zs.0 >= max_z,
+            "Falling piece should be higher than existing tower"
+        );
+
+        let new_y = max_z + zs.1 - zs.0 + 1;
+
+        for (x, y) in brick.iter_squares() {
+            let offset = x as usize + (y as usize) * width;
+            if height_map[offset] == max_z && max_z > 0 {
+                debug_assert_ne!(top_map[offset], usize::MAX);
+                temp.insert(top_map[offset]);
+            }
+            height_map[offset] = new_y;
+            top_map[offset] = i;
+        }
+
+        *supported = temp.len().try_into()?;
+        for support in temp.drain() {
+            supporting[support].push(i);
+        }
+    }
+
+    Ok((supporting, supported))
+}
+
+pub fn part1(input: &[u8]) -> anyhow::Result<String> {
+    let (supporting, supported) = compute_support(input)?;
+
+    let removable = supporting
+        .iter()
+        .filter(|b| b.iter().all(|&other| supported[other] >= 2))
+        .count();
+
+    Ok(removable.to_string())
+}
+
+fn compute_collapse(start: usize, supporting: &[Vec<usize>], supported: &[u8]) -> u32 {
+    // Need to make a copy to keep track of removed supports
+    let mut supported = supported.to_owned();
+    let mut todo = Vec::new();
+    todo.push(start);
+    let mut collapsed = 0;
+
+    while let Some(node) = todo.pop() {
+        for &other in &supporting[node] {
+            supported[other] -= 1;
+            if supported[other] == 0 {
+                todo.push(other);
+                collapsed += 1;
+            }
+        }
+    }
+
+    collapsed
+}
+
+pub fn part2(input: &[u8]) -> anyhow::Result<String> {
+    let (supporting, supported) = compute_support(input)?;
+
+    let collapsed = supporting
+        .iter()
+        .enumerate()
+        .filter(|(_, support)| support.iter().any(|&other| supported[other] == 1))
+        .map(|(start, _)| compute_collapse(start, &supporting, &supported))
+        .sum::<u32>();
+
+    Ok(collapsed.to_string())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const SAMPLE: &[u8] = include_bytes!("samples/22.txt");
+
+    #[test]
+    fn sample_part1() {
+        assert_eq!("5", part1(SAMPLE).unwrap());
+    }
+
+    #[test]
+    fn sample_part2() {
+        assert_eq!("7", part2(SAMPLE).unwrap());
+    }
 }

2023/src/samples/22.txt

@@ -0,0 +1,7 @@
+1,0,1~1,2,1
+0,0,2~2,0,2
+0,2,3~2,2,3
+0,0,4~0,2,4
+2,0,5~2,2,5
+0,1,6~2,1,6
+1,1,8~1,1,9