From 9e013b28016811c0b7f8956a03c3fd73cce8737b Mon Sep 17 00:00:00 2001
From: Greg Burri <greg.burri@gmail.com>
Date: Mon, 9 Dec 2024 15:34:10 +0100
Subject: [PATCH] Day[9], part 2

---
 src/day09.rs | 91 ++++++++++++++++++++++++++++++++++++++++++++--------
 src/days.rs  | 10 ++++--
 2 files changed, 85 insertions(+), 16 deletions(-)

diff --git a/src/day09.rs b/src/day09.rs
index 7dcd8e1..8ad49bb 100644
--- a/src/day09.rs
+++ b/src/day09.rs
@@ -1,33 +1,33 @@
 use std::io::BufRead;
 
+const EMPTY: u32 = u32::MAX;
+
 pub fn read(reader: &mut dyn BufRead) -> Vec<u32> {
     let mut disk_map_str = String::new();
     reader.read_to_string(&mut disk_map_str).unwrap();
-    disk_map_str
+    let disk_map: Vec<u32> = disk_map_str
         .chars()
         .map(|c| c.to_digit(10).unwrap())
-        .collect()
-}
-
-const EMPTY: u32 = u32::MAX;
+        .collect();
 
-pub fn defrag_checksum(disk_map: &[u32]) -> u64 {
-    let size: u32 = disk_map.iter().sum();
     let mut memory: Vec<u32> = Vec::new();
-    memory.resize(size as usize, EMPTY);
+    memory.resize(disk_map.iter().sum::<u32>() as usize, EMPTY);
 
     let mut current_id = 0;
     let mut current_pos: usize = 0;
     let mut is_free_space = false;
     for n in disk_map {
         if !is_free_space {
-            memory[current_pos..current_pos + *n as usize].fill(current_id);
+            memory[current_pos..current_pos + n as usize].fill(current_id);
             current_id += 1;
         }
-        current_pos += *n as usize;
+        current_pos += n as usize;
         is_free_space = !is_free_space;
     }
+    memory
+}
 
+pub fn defrag_v1(mut memory: Vec<u32>) -> Vec<u32> {
     fn next_empty_space_pos(from: usize, memory: &[u32]) -> usize {
         for i in from..memory.len() {
             if memory[i] == EMPTY {
@@ -56,10 +56,67 @@ pub fn defrag_checksum(disk_map: &[u32]) -> u64 {
         last_value_pos = prev_non_empty_space_pos(last_value_pos - 1, &memory);
     }
 
+    memory
+}
+
+struct MemoryChunk {
+    pos: usize,
+    size: usize,
+}
+
+pub fn defrag_v2(mut memory: Vec<u32>) -> Vec<u32> {
+    let mut free_spaces: Vec<MemoryChunk> = Vec::new();
+    let mut files: Vec<MemoryChunk> = Vec::new();
+
+    let mut pos = 0;
+    loop {
+        let start = pos;
+        let v = memory[start];
+        pos += 1;
+        while pos < memory.len() && memory[pos] == v {
+            pos += 1;
+        }
+
+        let chunk = MemoryChunk {
+            pos: start,
+            size: pos - start,
+        };
+
+        if v == EMPTY {
+            free_spaces.push(chunk);
+        } else {
+            files.push(chunk);
+        }
+
+        if pos == memory.len() {
+            break;
+        }
+    }
+
+    for f in files.iter().rev() {
+        for free_space in free_spaces.iter_mut() {
+            if free_space.pos >= f.pos {
+                break;
+            }
+            if free_space.size >= f.size {
+                let v = memory[f.pos];
+                memory[free_space.pos..free_space.pos + f.size].fill(v);
+                memory[f.pos..f.pos + f.size].fill(EMPTY);
+                free_space.size -= f.size;
+                free_space.pos += f.size;
+                break;
+            }
+        }
+    }
+
+    memory
+}
+
+pub fn checksum(memory: &[u32]) -> u64 {
     memory
         .iter()
-        .take_while(|v| **v != EMPTY)
         .enumerate()
+        .filter(|(_, v)| **v != EMPTY)
         .map(|(i, v)| i as u64 * *v as u64)
         .sum()
 }
@@ -72,11 +129,17 @@ mod tests {
 
     #[test]
     fn test_part1() {
-        let disk_map = read(&mut DISK_MAP.as_bytes());
-        let checksum = defrag_checksum(&disk_map);
+        let memory = read(&mut DISK_MAP.as_bytes());
+        let memory = defrag_v1(memory);
+        let checksum = checksum(&memory);
         assert_eq!(checksum, 1928);
     }
 
     #[test]
-    fn test_part2() {}
+    fn test_part2() {
+        let memory = read(&mut DISK_MAP.as_bytes());
+        let memory = defrag_v2(memory);
+        let checksum = checksum(&memory);
+        assert_eq!(checksum, 2858);
+    }
 }
diff --git a/src/days.rs b/src/days.rs
index 644051e..b6f3821 100644
--- a/src/days.rs
+++ b/src/days.rs
@@ -75,6 +75,12 @@ pub fn day08(reader: &mut dyn BufRead) -> String {
 }
 
 pub fn day09(reader: &mut dyn BufRead) -> String {
-    let disk_map = day09::read(reader);
-    format!("part1: {}, part2: {}", day09::defrag_checksum(&disk_map), 0,)
+    let memory = day09::read(reader);
+    let memory_defrag_v1 = day09::defrag_v1(memory.clone());
+    let memory_defrag_v2 = day09::defrag_v2(memory);
+    format!(
+        "part1: {}, part2: {}",
+        day09::checksum(&memory_defrag_v1),
+        day09::checksum(&memory_defrag_v2),
+    )
 }
-- 
2.45.2