AOC2024 -D11.rs

use std::{collections::HashMap, fs};

fn parse(input: &str) -> Vec<usize> {
    input
        .lines()
        .flat_map(|x| x.split_whitespace())
        .flat_map(|x| x.parse::<usize>())
        .collect::<Vec<usize>>()
}

fn split_maybe(stone: usize) -> Option<Vec<usize>> {
    let ilog10 = (stone.ilog10() + 1) % 2;
    if ilog10 % 2 != 0 {
        return None;
    }
    let exp = 10_usize.pow((stone.ilog10() + 1) / 2);
    let first_part = stone / exp;
    let second_part = stone % exp;

    Some(vec![first_part, second_part])
}

fn extends(stones: Vec<usize>) -> Vec<usize> {
    let mut new_stones = vec![];

    for stone in stones {
        if stone == 0 {
            new_stones.push(1);
        } else if let Some(mut splitted) = split_maybe(stone) {
            new_stones.append(&mut splitted);
        } else {
            new_stones.push(stone * 2024);
        }
    }
    new_stones
}

fn compute_recursive(stone: usize, blinks_to_do: usize) -> usize {
    if blinks_to_do == 0 {
        return 1;
    }

    if stone == 0 {
        compute_recursive(1, blinks_to_do - 1)
    } else if let Some(splitted) = split_maybe(stone) {
        splitted
            .iter()
            .map(|&x| compute_recursive(x, blinks_to_do - 1))
            .sum()
    } else {
        compute_recursive(stone * 2024, blinks_to_do - 1)
    }
}

fn map_count(map: &HashMap<usize, usize>) -> HashMap<usize, usize> {
    let mut new_map = HashMap::<usize, usize>::with_capacity(map.len() * 2);
    for (stone, count) in map {
        if *stone == 0 {
            *new_map.entry(1).or_default() += count;
        } else if let Some(splitted) = split_maybe(*stone) {
            *new_map.entry(splitted[0]).or_default() += count;
            *new_map.entry(splitted[1]).or_default() += count;
        } else {
            *new_map.entry(stone * 2024).or_default() += count;
        }
    }

    new_map
}

pub fn naive(input_path: &str, iteration: usize) -> anyhow::Result<usize> {
    let file = fs::read_to_string(input_path)?;

    let mut numbers = parse(&file);
    for _ in 0..iteration {
        numbers = extends(numbers);
    }

    Ok(numbers.len())
}

pub fn recursive(input_path: &str, iteration: usize) -> anyhow::Result<usize> {
    let file = fs::read_to_string(input_path)?;

    let numbers = parse(&file);
    let res = numbers
        .iter()
        .map(|&x| compute_recursive(x, iteration))
        .sum();
    Ok(res)
}

pub fn map(input_path: &str, iteration: usize) -> anyhow::Result<usize> {
    let file = fs::read_to_string(input_path)?;

    let mut map = HashMap::<usize, usize>::new();
    for stone in parse(&file) {
        *map.entry(stone).or_default() += 1;
    }

    for _ in 0..iteration {
        map = map_count(&map);
    }

    Ok(map.values().sum())
}


#[cfg(test)]
mod test {
    use crate::{map, naive, recursive, split_maybe};

    const ROOT_DIR: &str = env!("CARGO_MANIFEST_DIR");

    #[test]
    fn test_part1_naive() {
        assert_eq!(
            naive(&format!("{ROOT_DIR}/input_test.txt"), 25).unwrap(),
            55312
        );
    }
    #[test]
    fn test_part1_recursive() {
        assert_eq!(
            recursive(&format!("{ROOT_DIR}/input_test.txt"), 25).unwrap(),
            55312
        );
    }

    #[test]
    fn test_part1_map() {
        assert_eq!(
            map(&format!("{ROOT_DIR}/input_test.txt"), 25).unwrap(),
            55312
        );
    }

    #[test]
    fn test_part2() {
        assert_eq!(
            map(&format!("{ROOT_DIR}/input_test.txt"), 75).unwrap(),
            65601038650482
        );
    }

    #[test]
    fn test_split() {
        assert_eq!(split_maybe(1), None);
        assert_eq!(split_maybe(10), Some(vec![1, 0]));
        assert_eq!(split_maybe(100), None);
        assert_eq!(split_maybe(1000), Some(vec![10, 0]));
        assert_eq!(split_maybe(1500), Some(vec![15, 0]));
        assert_eq!(split_maybe(1060), Some(vec![10, 60]));
    }
}