AoC/aoc_2024/src/day7.rs

use aoc_runner_derive::{aoc, aoc_generator};
use nom::{
    bytes::complete::tag,
    character::complete::{newline, space1, u64 as parse_u64},
    error::Error,
    multi::separated_list1,
    sequence::separated_pair,
};

type Input = Vec<(u64, Vec<u64>)>;

#[aoc_generator(day7)]
fn parse(input: &str) -> Input {
    separated_list1(
        newline::<&str, Error<&str>>,
        separated_pair(parse_u64, tag(": "), separated_list1(space1, parse_u64)),
    )(input)
    .unwrap()
    .1
}

#[aoc(day7, part1)]
fn part1(input: &Input) -> usize {
    input
        .iter()
        .filter(|v| part1_test(v.1[0], v.0, &v.1[1..]))
        .fold(0, |acc, v| acc + v.0 as usize)
}

fn part1_test(curr: u64, target: u64, numbers: &[u64]) -> bool {
    if curr > target {
        false
    } else if numbers.is_empty() {
        curr == target
    } else {
        let num = numbers[0];
        let rest = &numbers[1..];
        part1_test(curr + num, target, rest) || part1_test(curr * num, target, rest)
    }
}

#[aoc(day7, part2)]
fn part2(input: &Input) -> usize {
    input
        .iter()
        .filter(|v| part2_test(v.1[0], v.0, &v.1[1..]))
        .fold(0, |acc, v| acc + v.0 as usize)
}

fn concat(a: u64, b: u64) -> u64 {
    let digits = b.ilog10() + 1;
    a * 10u64.pow(digits) + b
}

fn part2_test(curr: u64, target: u64, numbers: &[u64]) -> bool {
    if curr > target {
        false
    } else if numbers.is_empty() {
        curr == target
    } else {
        let num = numbers[0];
        let rest = &numbers[1..];
        part2_test(curr + num, target, rest)
            || part2_test(curr * num, target, rest)
            || part2_test(concat(curr, num), target, rest)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn part1_example() {
        assert_eq!(part1(&parse("190: 10 19\n3267: 81 40 27\n83: 17 5\n156: 15 6\n7290: 6 8 6 15\n161011: 16 10 13\n192: 17 8 14\n21037: 9 7 18 13\n292: 11 6 16 20")), 3749);
        assert_eq!(
            part1(&parse(include_str!("../input/2024/day7.txt"))),
            3598800864292
        );
    }

    #[test]
    fn part2_example() {
        assert_eq!(part2(&parse("190: 10 19\n3267: 81 40 27\n83: 17 5\n156: 15 6\n7290: 6 8 6 15\n161011: 16 10 13\n192: 17 8 14\n21037: 9 7 18 13\n292: 11 6 16 20")), 11387);
        assert_eq!(
            part2(&parse(include_str!("../input/2024/day7.txt"))),
            340362529351427
        );
    }
}