Advent of Code 2021 - Day 14: Extended Polymerization

extended_polymerization.rs

fn count(template: &[char], rules: &HashMap<(char, char), char>, steps: u64) -> u64 {
    let mut total = HashMap::new();
    let mut counts = HashMap::new();

    template.windows(2).for_each(|pair| {
        count_for_pair(pair[0], pair[1], rules, steps, &mut counts);
        merge(&mut total, &counts[&(pair[0], pair[1], steps)]);
    });

    template.iter().for_each(|&element| {
        *total.entry(element).or_insert(0) += 1;
    });

    total.values().max().unwrap() - total.values().min().unwrap()
}

fn count_for_pair(a: char, b: char, rules: &HashMap<(char, char), char>, steps: u64, counts: &mut HashMap<(char, char, u64), HashMap<char, u64>>) {
    if counts.get(&(a, b, steps)).is_some() {
        return;
    }

    let mut total = HashMap::new();

    let &inserted = rules.get(&(a, b)).unwrap();

    if steps > 1 {
        count_for_pair(a, inserted, rules, steps - 1, counts);
        count_for_pair(inserted, b, rules, steps - 1, counts);

        merge(&mut total, &counts[&(a, inserted, steps - 1)]);
        merge(&mut total, &counts[&(inserted, b, steps - 1)]);
    }

    *total.entry(inserted).or_insert(0) += 1;

    counts.insert((a, b, steps), total);
}

fn merge(into: &mut HashMap<char, u64>, from: &HashMap<char, u64>) {
    for c in from.keys() {
        *into.entry(*c).or_insert(0) += from[c];
    }
}