day7.rs

use pathfinding::prelude::topological_sort;
use regex::Regex;
use std::collections::{HashMap, HashSet};

fn parse_input(input: &str) -> HashMap<&str, Vec<(usize, &str)>> {
    let re1 = Regex::new(r"^(.*) bags contain (.*)$").unwrap();
    let re2 = Regex::new(r"(\d+) (.*?) bag").unwrap();
    input
        .lines()
        .map(|l| {
            (
                re1.captures(l).unwrap().get(1).unwrap().as_str(),
                re2.captures_iter(l)
                    .map(|c| (c[1].parse::<usize>().unwrap(), c.get(2).unwrap().as_str()))
                    .collect(),
            )
        })
        .collect()
}

#[aoc(day7, part1)]
fn part1(input: &str) -> usize {
    let mut reached_by = HashMap::new();
    for (n, v) in parse_input(input) {
        for (_, s) in v {
            reached_by
                .entry(s)
                .or_insert_with(|| HashSet::new())
                .insert(n);
        }
    }
    topological_sort(&["shiny gold"], |n| {
        reached_by.get(n).cloned().unwrap_or_else(HashSet::new)
    })
    .unwrap()
    .len()
        - 1
}

#[aoc(day7, part2)]
fn part2(input: &str) -> usize {
    let bags = parse_input(input);
    let sorted = topological_sort(&["shiny gold"], |c| bags[c].iter().map(|(_, d)| *d)).unwrap();
    let mut cap: HashMap<&str, usize> = HashMap::new();
    for (c, v) in sorted.into_iter().rev().map(|c| (c, &bags[c])) {
        cap.insert(c, v.iter().map(|&(n, ref d)| n * (1 + cap[d])).sum());
    }
    cap["shiny gold"]
}