advent-of-code/2023/src/bin/19.rs

advent_of_code::solution!(19);

use aoc_parse::{parser, prelude::*};
use std::collections::HashMap;

struct Part {
    x: u64,
    m: u64,
    a: u64,
    s: u64,
}

impl Part {
    fn sum(&self) -> u64 {
        self.x + self.m + self.a + self.s
    }

    fn val_for(&self, field: &str) -> u64 {
        match field {
            "x" => self.x,
            "m" => self.m,
            "a" => self.a,
            "s" => self.s,
            n @ _ => panic!("bad val_for field {}", n),
        }
    }
}

#[derive(Eq, PartialEq, Copy, Clone)]
enum RuleOp {
    NextState,
    LessThan,
    GreaterThan,
}

struct Rule {
    op: RuleOp,
    field: String,
    count: u64,
    dest: String,
}

impl Rule {
    fn next_state(&self) -> bool {
        self.op == RuleOp::NextState
    }

    fn match_lt(&self, part: &Part) -> bool {
        self.op == RuleOp::LessThan && part.val_for(&self.field) < self.count
    }

    fn match_gt(&self, part: &Part) -> bool {
        self.op == RuleOp::GreaterThan && part.val_for(&self.field) > self.count
    }
}

type Rules = HashMap<String, Vec<Rule>>;

fn parse(input: &str) -> (Rules, Vec<Part>) {
    let mut rules: Rules = HashMap::new();
    let rparse = parser!(string(lower+) "{" string(any_char+) "}");
    let ltparse = parser!(string(lower) "<" u64 ":" string(alpha+));
    let gtparse = parser!(string(lower) ">" u64 ":" string(alpha+));
    for line in input.lines() {
        if line.is_empty() {
            break;
        }
        let (rname, raw_rules) = rparse.parse(line).unwrap();
        let new_rules = raw_rules
            .split(",")
            .map(|t| {
                if t.contains("<") {
                    let (f, c, d) = ltparse.parse(t).unwrap();
                    Rule {
                        op: RuleOp::LessThan,
                        field: f,
                        count: c,
                        dest: d,
                    }
                } else if t.contains(">") {
                    let (f, c, d) = gtparse.parse(t).unwrap();
                    Rule {
                        op: RuleOp::GreaterThan,
                        field: f,
                        count: c,
                        dest: d,
                    }
                } else {
                    Rule {
                        op: RuleOp::NextState,
                        field: String::from("f"),
                        count: 0,
                        dest: String::from(t),
                    }
                }
            })
            .collect::<Vec<_>>();
        rules.insert(rname, new_rules);
    }

    let mut parts: Vec<Part> = vec![];
    let pparse = parser!("{x=" u64 ",m=" u64 ",a=" u64 ",s=" u64 "}");
    let mut skipped = false;
    for line in input.lines() {
        if !skipped {
            if line.is_empty() {
                skipped = true;
            }
            continue;
        }
        let (x, m, a, s) = pparse.parse(line).unwrap();
        parts.push(Part { x, m, a, s });
    }

    (rules, parts)
}

fn filter_part(rules: &Rules, part: &Part) -> u64 {
    let mut b = rules.get("in").unwrap();
    loop {
        let mut found = false;
        for r in b.iter() {
            if r.match_lt(part) || r.match_gt(part) || r.next_state() {
                match r.dest.as_str() {
                    "A" => return part.sum(),
                    "R" => return 0,
                    _ => {
                        found = true;
                        b = rules.get(&r.dest).unwrap();
                        break;
                    }
                }
            }
        }
        if !found {
            break;
        }
    }
    0
}

pub fn part_one(input: &str) -> Option<u64> {
    let (rules, parts) = parse(input);
    Some(parts.iter().map(|p| filter_part(&rules, &p)).sum())
}

fn combinations(weights: [(u64, u64); 4]) -> u64 {
    (0..4)
        .into_iter()
        .map(|i| weights[i].0.abs_diff(weights[i].1))
        .product()
}

fn windex(field: &str) -> usize {
    match field {
        "x" => 0,
        "m" => 1,
        "a" => 2,
        "s" => 3,
        _ => panic!("bad field"),
    }
}

fn split_lt(count: u64, input: (u64, u64)) -> ((u64, u64), (u64, u64)) {
    let one = input.0..input.1.min(count);
    let two = count..input.1;

    ((one.start, one.end), (two.start, two.end))
}

fn split_gt(count: u64, input: (u64, u64)) -> ((u64, u64), (u64, u64)) {
    let one = (count + 1)..input.1;
    let two = input.0..(count + 1);

    ((one.start, one.end), (two.start, two.end))
}

fn traverse_rules(rules: &Rules, cur: &str, mut weights: [(u64, u64); 4]) -> u64 {
    let mut c = 0;
    for r in rules.get(cur).unwrap().iter() {
        match r.op {
            RuleOp::NextState => {
                c += match r.dest.as_str() {
                    "A" => combinations(weights),
                    "R" => 0,
                    _ => traverse_rules(rules, &r.dest, weights),
                }
            }
            RuleOp::LessThan => {
                let i = windex(&r.field);
                let mut next_weights = weights.clone();
                let (next_range, cur_range) = split_lt(r.count, weights[i]);
                weights[i] = cur_range;
                next_weights[i] = next_range;
                c += match r.dest.as_str() {
                    "A" => combinations(next_weights),
                    "R" => 0,
                    _ => traverse_rules(rules, &r.dest, next_weights),
                }
            }
            RuleOp::GreaterThan => {
                let i = windex(&r.field);
                let mut next_weights = weights.clone();
                let (next_range, cur_range) = split_gt(r.count, weights[i]);
                weights[i] = cur_range;
                next_weights[i] = next_range;
                c += match r.dest.as_str() {
                    "A" => combinations(next_weights),
                    "R" => 0,
                    _ => traverse_rules(rules, &r.dest, next_weights),
                }
            }
        }
    }
    c
}

pub fn part_two(input: &str) -> Option<u64> {
    let (rules, _parts) = parse(input);
    Some(traverse_rules(&rules, "in", [(1, 4001); 4]))
}

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

    #[test]
    fn test_part_one() {
        let result = part_one(&advent_of_code::template::read_file("examples", DAY));
        assert_eq!(result, Some(19114));
    }

    #[test]
    fn test_part_two() {
        let result = part_two(&advent_of_code::template::read_file("examples", DAY));
        assert_eq!(result, Some(167409079868000));
    }
}