From e1ecbe536cc06fb8906ee1ebe757c49844844cfe Mon Sep 17 00:00:00 2001
From: uku <hi@uku.moe>
Date: Tue, 10 Dec 2024 01:33:53 +0100
Subject: [PATCH] feat(day9): init

---
 src/solutions/day_09.rs | 114 ++++++++++++++++++++++++++++++++++++++++
 src/solutions/mod.rs    |   2 +
 2 files changed, 116 insertions(+)
 create mode 100644 src/solutions/day_09.rs

diff --git a/src/solutions/day_09.rs b/src/solutions/day_09.rs
new file mode 100644
index 0000000..76f3d93
--- /dev/null
+++ b/src/solutions/day_09.rs
@@ -0,0 +1,114 @@
+use itertools::Itertools;
+
+use crate::common::{Answer, Solution};
+
+pub struct Day09;
+
+impl Solution for Day09 {
+    fn name(&self) -> &'static str {
+        "Disk Fragmenter"
+    }
+
+    fn part_a(&self, input: &str) -> Answer {
+        let nums = input
+            .chars()
+            .map(|c| c.to_digit(10).unwrap())
+            .collect::<Vec<_>>();
+
+        let mut aligned_files = nums
+            .iter()
+            .step_by(2)
+            .enumerate()
+            .flat_map(|(i, &n)| vec![i; n as usize])
+            .collect::<Vec<_>>();
+
+        let mut final_disk = Vec::new();
+
+        for (i, n) in nums.iter().enumerate() {
+            if aligned_files.is_empty() {
+                break;
+            }
+
+            let len = (*n as usize).min(aligned_files.len());
+
+            if i % 2 == 0 {
+                final_disk.extend(aligned_files.drain(0..len));
+            } else {
+                final_disk.extend(
+                    aligned_files
+                        .split_off(aligned_files.len() - len)
+                        .iter()
+                        .rev(),
+                );
+            };
+        }
+
+        final_disk
+            .into_iter()
+            .enumerate()
+            .map(|(i, n)| i * n)
+            .sum::<usize>()
+            .into()
+    }
+
+    fn part_b(&self, input: &str) -> Answer {
+        let mut chunks = input
+            .chars()
+            .map(|c| c.to_digit(10).unwrap())
+            .enumerate()
+            .map(|(i, n)| ((i % 2 == 0).then_some(i >> 1), n))
+            .collect::<Vec<_>>();
+
+        let files = chunks
+            .iter()
+            .filter_map(|(o, n)| o.map(|i| (i, *n)))
+            .rev()
+            .collect::<Vec<_>>();
+
+        for (content, size) in files {
+            let free_idx = chunks
+                .iter()
+                .position(|(o, len)| o.is_none() && *len >= size);
+
+            // we need to recompute this because stuff moves in the vec and the element count is not constant
+            let orig_idx = chunks
+                .iter()
+                .position(|chunk| chunk == &(Some(content), size))
+                .unwrap();
+
+            if let Some(free_idx) = free_idx {
+                if free_idx < orig_idx {
+                    chunks.get_mut(orig_idx).unwrap().0 = None;
+                    chunks.insert(free_idx, (Some(content), size));
+                    chunks.get_mut(free_idx + 1).unwrap().1 -= size;
+                }
+            }
+        }
+
+        chunks
+            .iter()
+            .flat_map(|(o, n)| vec![o.unwrap_or(0); *n as usize])
+            .enumerate()
+            .map(|(i, n)| i * n)
+            .sum::<usize>()
+            .into()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::Day09;
+    use crate::common::Solution;
+
+    const INPUT: &str = "2333133121414131402";
+
+    #[test]
+    fn part_a() {
+        assert_eq!(Day09.part_a(INPUT), 1928.into());
+    }
+
+    #[test]
+    fn part_b() {
+        assert_eq!(Day09.part_b(INPUT), 2858.into());
+    }
+}
diff --git a/src/solutions/mod.rs b/src/solutions/mod.rs
index 798dee2..c5c69cf 100644
--- a/src/solutions/mod.rs
+++ b/src/solutions/mod.rs
@@ -8,6 +8,7 @@ mod day_05;
 mod day_06;
 mod day_07;
 mod day_08;
+mod day_09;
 
 pub const SOLUTIONS: &[&dyn Solution] = &[
     &day_01::Day01,
@@ -18,4 +19,5 @@ pub const SOLUTIONS: &[&dyn Solution] = &[
     &day_06::Day06,
     &day_07::Day07,
     &day_08::Day08,
+    &day_09::Day09,
 ];