diff --git a/src/day9.rs b/src/day9.rs index fe0e366..0a39f1c 100644 --- a/src/day9.rs +++ b/src/day9.rs @@ -2,17 +2,98 @@ pub fn day9_part1(input: &str) -> String { let diskmap = parse(input); let mut blocks = diskmap_to_blocks(&diskmap); compact(&mut blocks); - let sum = blocks - .into_iter() - .flatten() - .enumerate() - .map(|(index, block)| index as u64 * block) - .sum::(); + let sum = checksum(&blocks); sum.to_string() } pub fn day9_part2(input: &str) -> String { - todo!() + let diskmap = parse(input); + let compacted_diskmap = defrag_compact(&diskmap); + let blocks = compacted_diskmap_to_blocks(&compacted_diskmap); + let sum = checksum(&blocks); + sum.to_string() +} + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +struct File { + id: u64, + length: usize, + free_space: usize, +} + +fn compacted_diskmap_to_blocks(compacted_diskmap: &[File]) -> Vec> { + compacted_diskmap + .iter() + .flat_map(|file| { + [ + vec![Some(file.id); file.length], + vec![None; file.free_space], + ] + }) + .flatten() + .collect() +} + +fn defrag_compact(diskmap: &[usize]) -> Vec { + let files = diskmap_to_files(diskmap); + let mut compacted_files = files.clone(); + + for file in files.clone().into_iter().rev() { + let Some(insertable_index) = compacted_files + .iter() + .enumerate() + .find(|(_, insertable_file)| insertable_file.free_space >= file.length) + .map(|(i, _)| i) + else { + continue; + }; + + let prev_free_space = compacted_files[insertable_index].free_space; + compacted_files[insertable_index].free_space = 0; + compacted_files.insert(insertable_index + 1, file); + compacted_files[insertable_index + 1].free_space = prev_free_space - file.length; + + let nullable_index = compacted_files + .iter() + .enumerate() + .rev() + .find(|(_, nullable_file)| nullable_file.id == file.id) + .map(|(i, _)| i) + .unwrap(); + compacted_files[nullable_index - 1].free_space += + compacted_files[nullable_index].length + compacted_files[nullable_index].free_space; + compacted_files.remove(nullable_index); + } + + compacted_files +} + +fn diskmap_to_files(diskmap: &[usize]) -> Vec { + let mut id = 0u64; + let mut blocks = vec![]; + + for digit in diskmap { + if id % 2 == 0 { + blocks.push(File { + id: id / 2, + length: *digit, + free_space: 0, + }); + } else { + blocks.last_mut().unwrap().free_space = *digit; + } + id += 1 + } + + blocks +} + +fn checksum(blocks: &[Option]) -> u64 { + blocks + .iter() + .enumerate() + .map(|(index, block)| if let Some(block) = block {index as u64 * block} else {0}) + .sum::() } fn diskmap_to_blocks(diskmap: &[usize]) -> Vec> { diff --git a/src/lib.rs b/src/lib.rs index 8b9d499..9bd26be 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -357,13 +357,13 @@ mod tests { let result = day9::day9_part1(include_str!("../input/day9.txt")); assert_eq!(result, "6448989155953"); } - // #[test] - // fn day9_part2_test() { - // let test_result = day9::day9_part2(include_str!("../input/day9.test.txt")); - // assert_eq!(test_result, "34"); - // // let result = day9::day9_part2(include_str!("../input/day9.txt")); - // // assert_eq!(result, "1169"); - // } + #[test] + fn day9_part2_test() { + let test_result = day9::day9_part2(include_str!("../input/day9.test.txt")); + assert_eq!(test_result, "2858"); + let result = day9::day9_part2(include_str!("../input/day9.txt")); + assert_eq!(result, "6476642796832"); + } #[test] fn day8_part1_test() {