advent-of-code-2024-rust/src/day9.rs

use crate::day_solver::DaySolver;
use num::Integer;

pub struct Day9 {
    disk_map: Vec<u8>
}

impl Day9 {

    pub fn create(input: String) -> Self {
        Day9 {
            disk_map: input.chars().map(|c| (c as u8) - b'0').collect()
        }
    }

    // fn print_layout(disk_blocks: &Vec<Block>) {
    //     for block in disk_blocks {
    //         match block {
    //             Block::File(size, idx) => {
    //                 for _ in 0..*size {
    //                     print!("{}", idx);
    //                 }
    //             }
    //             Block::Empty(size) => {
    //                 print!("{}", ".".repeat(*size));
    //             }
    //         }
    //     }
    //     println!();
    // }
}

impl DaySolver for Day9 {


    fn solve_part1(&mut self) -> String {

        let mut cur_disk_idx = 0usize;
        let mut disk_map_idx = 0usize;
        let mut disk_map_idx_rev = self.disk_map.len() + (self.disk_map.len() % 2);
        let mut left_to_move = 0u8;

        let mut res = 0usize;

        while disk_map_idx <= disk_map_idx_rev {
            if disk_map_idx.is_even() {
                // We're reading a file from the front of the disk map
                let left_in_file = if disk_map_idx == disk_map_idx_rev { left_to_move } else { self.disk_map[disk_map_idx] };
                for _ in 0..left_in_file {
                    // print!("{}({})", disk_map_idx / 2, cur_disk_idx);
                    res += cur_disk_idx * (disk_map_idx / 2);
                    cur_disk_idx += 1;
                }
            } else {
                // We're reading empty space
                for _ in 0..self.disk_map[disk_map_idx] {
                    // So let's grab some blocks from files at the end of the disk map
                    while left_to_move == 0 &&  disk_map_idx_rev > disk_map_idx {
                        disk_map_idx_rev -= 2;
                        left_to_move = self.disk_map[disk_map_idx_rev];
                    }
                    if disk_map_idx_rev < disk_map_idx {
                        break
                    }
                    // print!("{}({})", disk_map_idx_rev / 2, cur_disk_idx);
                    res += cur_disk_idx * (disk_map_idx_rev / 2);
                    cur_disk_idx += 1;
                    left_to_move -= 1;
                }
            }
            disk_map_idx += 1;
        }

        res.to_string()
    }

    fn solve_part2(&mut self) -> String {

        // Somewhat nicer data structure
        let mut disk_blocks: Vec<Block> = Vec::new();
        for (i, b) in self.disk_map.iter().enumerate() {
            if i.is_even() {
                disk_blocks.push(Block::File(b.to_owned().into(), i / 2))
            } else if b > &0 {
                disk_blocks.push(Block::Empty(b.to_owned().into()))
            }
        }

        // Self::print_layout(&disk_blocks);

        let original_disk = disk_blocks.clone();
        for block in original_disk.iter().rev() {
            if let Block::File(size, idx) = block {
                // Self::print_layout(&disk_blocks);
                // println!("{:?}", disk_blocks);
                // Try to move the file
                let move_target = disk_blocks.iter().enumerate()
                    .find(|(_, b) | if let Block::Empty(s) = b { s >= size } else { false })
                    // To prevent double borrowing, we need to clone the block here
                    .map(|(i, b)| (i, b.to_owned()));
                if let Some((can_move_to_idx, b)) = move_target.to_owned() {
                    match b {
                        Block::Empty(empty_size) => {
                            let original_index = disk_blocks.iter().position(|b| b == block).unwrap();
                            // Never move blocks to the right
                            if can_move_to_idx >= original_index { continue; }

                            // We replace the original block by an empty one
                            // Note that, because we never move blocks to the right, it doesn't matter that the empty space around the block we're moving is inaccurate.
                            disk_blocks[original_index] = Block::Empty(size.to_owned());
                            disk_blocks[can_move_to_idx] = Block::File(size.to_owned(), idx.to_owned());
                            if &empty_size > size {
                                // Add additional empty space
                                disk_blocks.insert(can_move_to_idx + 1, Block::Empty(empty_size.to_owned() - size.to_owned()))
                            }
                        }
                        _ => panic!("uhm?")
                    }
                }
            }
        }

        // Self::print_layout(&disk_blocks);
        let mut checksum = 0usize;
        let mut cur_disk_idx = 0usize;
        for block in disk_blocks {
            match block {
                Block::File(size, idx) => {
                    for _ in 0..size {
                        checksum += cur_disk_idx * idx;
                        cur_disk_idx += 1;
                    }
                }
                Block::Empty(size) => {
                    cur_disk_idx += size;
                }
            }
        }
        checksum.to_string()
    }
}

#[derive(Copy, Clone, PartialEq, Debug)]
enum Block {
    Empty(usize),
    File(usize, usize)
}

#[test]
fn test_part1_basic() {
    let mut day = Day9::create("12345".to_string());
    assert_eq!("60", day.solve_part1());
}

#[test]
fn test_part1() {
    let mut day = Day9::create("2333133121414131402".to_string());
    assert_eq!("1928", day.solve_part1());
}

#[test]
fn test_part2() {
    let mut day = Day9::create("2333133121414131402".to_string());
    assert_eq!("2858", day.solve_part2());
}

#[test]
fn test_part2_bonus() {
    let mut day = Day9::create("2931514".to_string());
    assert_eq!("158", day.solve_part2());
}