[2024] Slight day 9 cleanup

aoc_2024/src/day_09.rs

@@ -1,76 +1,94 @@
 use common::{solution, Answer};
-use itertools::Itertools;
+use itertools::{repeat_n, Itertools};
 
 solution!("Disk Fragmenter", 9);
 
 fn part_a(input: &str) -> Answer {
-    let mut problem = Problem::parse(input);
-    problem.sort_blocks();
-    problem.score_blocks().into()
+    let mut problem = Blocks::parse(input);
+    problem.sort();
+    problem.score().into()
 }
 
 fn part_b(input: &str) -> Answer {
-    let mut problem = Problem::parse(input);
-    problem.sort_files();
-    problem.score_files().into()
+    let mut problem = Files::parse(input);
+    problem.sort();
+    problem.score().into()
 }
 
-#[derive(Debug)]
-struct Problem {
+struct Blocks {
     blocks: Vec<Option<u32>>,
+}
+
+struct Files {
     files: Vec<File>,
 }
 
-#[derive(Debug)]
 struct File {
     pos: u32,
     size: u8,
     id: u32,
 }
 
-impl Problem {
+impl Blocks {
     fn parse(input: &str) -> Self {
         let mut blocks = Vec::new();
-        let mut files = Vec::new();
 
         let mut id = 0;
-        let mut pos = 0;
-
         for (idx, chr) in input.trim().char_indices() {
-            let num = chr.to_digit(10).unwrap() as u8;
+            let count = chr.to_digit(10).unwrap() as u8;
 
-            if idx % 2 == 1 {
-                for _ in 0..num {
-                    blocks.push(None);
-                }
-            } else {
-                files.push(File { pos, size: num, id });
-                for _ in 0..num {
-                    blocks.push(Some(id));
-                }
-                id += 1;
-            }
+            let is_block = idx % 2 == 0;
+            let item = is_block.then_some(id);
 
-            pos += num as u32;
+            blocks.extend(repeat_n(item, count as usize));
+            id += is_block as u32;
         }
 
-        Self { blocks, files }
+        Self { blocks }
     }
 
-    fn sort_blocks(&mut self) {
+    fn sort(&mut self) {
         loop {
             let empty = self.blocks.iter().position(|x| x.is_none()).unwrap();
             let last = self.blocks.iter().rposition(|x| x.is_some()).unwrap();
 
-            if last > empty {
-                self.blocks.swap(empty, last);
-            } else {
+            if last <= empty {
                 break;
             }
+
+            self.blocks.swap(empty, last);
         }
     }
 
-    fn sort_files(&mut self) {
+    fn score(&self) -> u64 {
+        self.blocks
+            .iter()
+            .enumerate()
+            .map(|(idx, id)| idx as u64 * id.unwrap_or_default() as u64)
+            .sum()
+    }
+}
+
+impl Files {
+    fn parse(input: &str) -> Self {
+        let mut files = Vec::new();
+        let (mut id, mut pos) = (0, 0);
+
+        for (idx, chr) in input.trim().char_indices() {
+            let size = chr.to_digit(10).unwrap() as u8;
+
+            if idx % 2 == 0 {
+                files.push(File { pos, size, id });
+                id += 1;
+            }
+
+            pos += size as u32;
+        }
+
+        Self { files }
+    }
+
+    fn sort(&mut self) {
         let max_id = self.files.last().unwrap().id;
         for id in (0..=max_id).rev() {
             let file_idx = self.files.iter().position(|x| x.id == id).unwrap();
@@ -99,32 +117,20 @@ impl Problem {
         }
     }
 
-    fn score_blocks(&self) -> u64 {
-        let mut sum = 0;
-
-        for (idx, id) in self.blocks.iter().enumerate() {
-            sum += idx as u64 * id.unwrap_or_default() as u64;
-        }
-
-        sum
-    }
-
-    fn score_files(&self) -> u64 {
+    fn score(&self) -> u64 {
         let mut sum = 0;
         let mut last = 0;
         let mut idx = 0;
+
         for x in &self.files {
-            for _ in last..x.pos {
-                idx += 1;
-            }
+            idx += x.pos - last;
 
-            for _ in 0..x.size {
-                sum += idx * x.id as u64;
-                idx += 1;
-            }
+            sum += (x.id as u64 * x.size as u64 * (x.size as u64 + 2 * idx as u64 - 1)) / 2;
+            idx += x.size as u32;
 
             last = x.pos + x.size as u32;
         }
+
         sum
     }
 }