day21.rs

use itertools::Itertools;
use pathfinding::prelude::{kuhn_munkres, Matrix};
use std::collections::{hash_map::Entry, BTreeSet, HashMap, HashSet};

type Food = (HashSet<String>, HashSet<String>);

#[aoc_generator(day21)]
fn input_generator(input: &str) -> Vec<Food> {
    input
        .lines()
        .map(|l| {
            let mut l = l.split("(contains ");
            (
                l.next()
                    .unwrap()
                    .split_ascii_whitespace()
                    .map(|s| s.to_owned())
                    .collect(),
                {
                    let a = l.next().unwrap();
                    a[..a.len() - 1].split(", ").map(|s| s.to_owned()).collect()
                },
            )
        })
        .collect()
}

#[aoc(day21, part1)]
fn part1(food: &[Food]) -> usize {
    let maybe_in = maybe_in(food);
    let all_ing = food
        .iter()
        .flat_map(|(i, _)| i.iter())
        .collect::<HashSet<_>>();
    let pot = maybe_in.values().flat_map(|s| s.iter().cloned()).collect();
    let safe: HashSet<&String> = all_ing.difference(&pot).cloned().collect();
    food.iter()
        .map(|(i, _)| safe.intersection(&i.iter().collect()).count())
        .sum()
}

fn maybe_in(food: &[Food]) -> HashMap<&String, HashSet<&String>> {
    let mut maybe_in = HashMap::new();
    for (ing, all) in food.iter() {
        for a in all.iter() {
            match maybe_in.entry(a) {
                Entry::Vacant(e) => {
                    e.insert(ing.iter().collect::<HashSet<_>>());
                }
                Entry::Occupied(mut e) => {
                    let v = e.get_mut();
                    *v = v.intersection(&ing.iter().collect()).cloned().collect();
                }
            }
        }
    }
    maybe_in
}

#[aoc(day21, part2)]
fn part2(food: &[Food]) -> String {
    let maybe_in = maybe_in(food);
    let all_all = food
        .iter()
        .flat_map(|(_, a)| a.iter())
        .collect::<BTreeSet<_>>();
    let pot = maybe_in
        .values()
        .flat_map(|s| s.iter().cloned())
        .unique()
        .collect::<Vec<_>>();
    kuhn_munkres(
        &Matrix::from_rows(all_all.iter().map(|&a| {
            let m = &maybe_in[&a];
            pot.iter().map(move |ing| m.contains(ing) as i32)
        }))
        .unwrap(),
    )
    .1
    .into_iter()
    .map(|i| pot[i].as_str())
    .collect::<Vec<_>>()
    .join(",")
}