2024_day7.rs

use std::env;
use std::fs;

use rayon::prelude::*;

fn get_arg() -> String {
    match env::args().nth(1) {
        Some(f) => f,
        None => "example".into(),
    }
}

fn get_content(filename: String) -> String {
    fs::read_to_string(filename).unwrap().trim().to_string()
}

fn parse(s: String) -> Vec<Equation> {
    s.split('\n').map(Equation::parse).collect()
}

#[derive(Debug)]
struct Equation {
    result: i64,
    numbers: Vec<i64>,
}

impl Equation {
    fn parse(line: &str) -> Self {
        let (result, rest) = line.split_once(':').unwrap();
        let numbers = rest
            .split_ascii_whitespace()
            .map(|n| n.parse().unwrap())
            .collect();

        Equation{
            result: result.parse().unwrap(),
            numbers,
        }
    }

    fn has_solution(&self) -> bool {
        self.search_solution(1, self.numbers[0])
    }

    fn search_solution(&self, index: usize, accumulator: i64) -> bool {
        if index == self.numbers.len() {
            return accumulator == self.result;
        }

        let n = self.numbers[index];
        self.search_solution(index + 1, n + accumulator) ||
            self.search_solution(index + 1, n * accumulator)
    }

    fn has_solution2(&self) -> bool {
        self.search_solution2(1, self.numbers[0])
    }

    fn search_solution2(&self, index: usize, accumulator: i64) -> bool {
        if index == self.numbers.len() {
            return accumulator == self.result;
        }

        let n = self.numbers[index];
        self.search_solution2(index + 1, n + accumulator) ||
            self.search_solution2(index + 1, n * accumulator) ||
            self.search_solution2(index + 1, concat(accumulator, n))
    }
}

#[inline(always)]
fn concat(left: i64, right: i64) -> i64 {
    let mut shift = 1;
    while shift <= right {
        shift *= 10;
    }

    (left * shift) + right
}

fn part1(equations: &[Equation]) -> i64 {
    equations.iter()
        .filter(|e| e.has_solution())
        .map(|e| e.result)
        .sum()
}

fn part2(equations: &[Equation]) -> i64 {
    equations.iter()
        .filter(|e| e.has_solution2())
        .map(|e| e.result)
        .sum()
}

fn search_solution(result: i64, accum: i64, numbers: &[i64]) -> bool {
    if let Some((n, rest)) = numbers.split_first() {
        search_solution(result, accum + n, rest) ||
            search_solution(result, accum * n, rest) ||
            search_solution(result, concat(accum, *n), rest)
    } else {
        result == accum
    }
}

fn part2_fast(equations: &[Equation]) -> i64 {
    equations
        .par_iter()
        .map(|eq| {
            let (acc, nums) = eq.numbers.split_first().unwrap();
            if search_solution(eq.result, *acc, nums) {
                eq.result
            } else {
                0
            }
        })
        .sum()
}

fn main() {
    let equations = parse(get_content(get_arg()));

    use std::time::Instant;
    let start = Instant::now();
    println!("Part 1: {}", part1(&equations));
    let part1_time = start.elapsed();

    let start2 = Instant::now();
    println!("Part 2: {}", part2(&equations));
    let part2_time = start2.elapsed();

    let start3 = Instant::now();
    println!("Part 2 fast: {}", part2_fast(&equations));
    let fast_time = start3.elapsed();

    println!("Part 1 time: {:?}", part1_time);
    println!("Part 2 time: {:?}", part2_time);
    println!("Part 2 fast time: {:?}", fast_time);
}