Day09 - part 2

diff --git a/src/day09.rs b/src/day09.rs
index 42a1409ce4cc58db2cbd38691abc243f312f54bf..ac087a3697754c7b4659b0806c3cf07120ae7dd1 100644 (file)
--- a/src/day09.rs
+++ b/src/day09.rs
@@ -49,9 +49,9 @@ impl Puzzle {
         block_idx
     }
 
-    fn first_free_idx(self: &Self) -> Option<usize> {
+    fn first_free_idx(self: &Self, size: u64) -> Option<usize> {
         self.free.iter().enumerate()
-            .filter(|(_, (_, length))| *length > 0)
+            .filter(|(_, (_, length))| *length >= size)
             .map(|(i, _)| i)
             .nth(0)
     }
@@ -67,7 +67,7 @@ impl Puzzle {
     fn compact(self: &mut Self) {
         let mut block_idx = 0;
         while self.can_compact() {
-            let first_free = self.first_free_idx();
+            let first_free = self.first_free_idx(1);
             if first_free.is_some() {
                 let free_idx = first_free.unwrap();
                 block_idx = self.fill_free(free_idx, block_idx);
@@ -77,6 +77,31 @@ impl Puzzle {
             }
         }
     }
+
+    fn compact_by_size(self: &mut Self) {
+        for i in (0..self.blocks.len()).rev() {
+            match self.first_free_idx(self.blocks[i].1) {
+                Some(first_free_idx) => {
+                    let first_free = self.free.get_mut(first_free_idx).unwrap();
+                    let block = self.blocks.get_mut(i).unwrap();
+                    if block.0 > first_free.0 {
+                        block.0 = first_free.0;
+                        first_free.0 += block.1;
+                        first_free.1 -= block.1;
+                    }
+                },
+                _ => {}
+            }
+        }
+    }
+
+    fn checksum(self: &Self) -> u64 {
+        self.blocks.iter()
+            .map(|&(start, length, value)| {
+                (start..(start+length)).map(|i| i*value).sum::<u64>()
+            })
+            .sum()
+    }
 }
 
 fn run_part1(input: &str) -> Result<u64, Box<dyn Error>> {
@@ -84,19 +109,15 @@ fn run_part1(input: &str) -> Result<u64, Box<dyn Error>> {
 
     let mut puzzle = Puzzle::new(input);
     puzzle.compact();
-    let res = puzzle.blocks.into_iter()
-        .map(|(start, length, value)| {
-            (start..(start+length)).map(|i| i*value).sum::<u64>()
-        })
-        .sum();
-
-    Ok(res)
+    Ok(puzzle.checksum())
 }
 
 fn run_part2(input: &str) -> Result<u64, Box<dyn Error>> {
     println!("Running {} - part 2", get_day());
 
-    Ok(0)
+    let mut puzzle = Puzzle::new(input);
+    puzzle.compact_by_size();
+    Ok(puzzle.checksum())
 }
 
 pub fn run(input: &str) -> Result<(), Box<dyn Error>> {
@@ -131,6 +152,7 @@ mod tests {
 
     #[test]
     fn test_part2() {
-        assert_eq!(0, run_part2(TEXT_INPUT).unwrap());
+        assert_eq!(132, run_part2(TEXT_INPUT_0).unwrap());
+        assert_eq!(2858, run_part2(TEXT_INPUT).unwrap());
     }
 }