Advent of Code 2020 Solution Megathread - Day 21: Allergen Assessment

import re
from collections import defaultdict
rule = re.compile("^((?:\w+ )+)\(contains (.*?)\)$")

data = open("demo.txt").read().splitlines()
all_candidates = defaultdict(list)
all_ingredients = []
for line in data:
    ing_str, all_str = rule.match(line).groups()

    ingredients = ing_str.strip().split(" ")
    allergens = all_str.split(", ")

    for allergen in allergens:
        all_candidates[allergen].append(set(ingredients))

    all_ingredients.extend(ingredients)

candidates = {}
for allergen_name, ingredient_list in all_candidates.items():
    candidates[allergen_name] = set.intersection(*ingredient_list)

print("appearances", sum(all_ingredients.count(ingredient) for ingredient in set(all_ingredients).difference(set.union(*candidates.values()))))

confirmed = {}

while len(confirmed) < len(candidates):
    for allergen in candidates:
        new_set = candidates[allergen].difference(confirmed.values())

        if len(new_set) == 1:
            confirmed[allergen] = new_set.pop()

        candidates[allergen] = new_set

print(",".join(v for k, v in sorted(confirmed.items(), key=lambda item: item[0])))

split :: String -> Char -> String -> [String]
split [] _ ss | null ss = []
              | otherwise = [ss]
split (x:xs) c ss | c == x = ss : split xs c ""
                  | otherwise = split xs c (ss ++ [x])

parse :: [String] -> [(S.Set String, S.Set String)]
parse = map parse'
    where
        parse' xs = let [fs, as] = split xs '(' ""
                        fs'      = split fs ' ' ""
                        (a:as')  = map (takeWhile (/=')') . dropWhile (==' ')) (split as ',' "")
                    in (S.fromList fs', S.fromList $ drop 9 a:as')

generate foods = valid <$> 
                 M.fromListWith (++) (do (is, as) <- foods
                                         a       <- S.toList as
                                         return (a, [is]))
    where                               
        valid is = S.filter (\g -> all (g `S.member`) is) $ S.unions is

task1 foods = length $ filter (`S.notMember` allValids) $ concatMap (S.toList . fst) foods
 where
  isValid ing = any (\a -> ing `S.member` (generate foods M.! a))
  allValids = S.unions $ map (\(ings, a) -> S.filter (`isValid` a) ings) foods

task2 = intercalate "," . map (S.elemAt 0 . snd) . M.toAscList . go . generate
   where
    -- Just like day 16 find foods which contain a particular allergen and repeat
    -- though a process of elimination
    go :: M.Map String (S.Set String) -> M.Map String (S.Set String)
    go v
      | all (\s -> S.size s == 1) (M.elems v) = v
      | otherwise = let found = S.unions $ M.elems $ M.filter (\s -> S.size s == 1) v
                    in go $ M.map (\s -> if S.size s == 1 then s else s S.\\ found) v

main = do
    is <- readFile "day21_input" <&> lines <&> parse
    print (task1 is)
    print (task2 is)

"""Day 21: Allergen Assessment

Find out through process of elimination which foreign ingredients
contain which allergens.
"""

from collections import Counter


def parse(filename):
    """Parse the input file into Python objects.

    Input file has lines like this:
        sdfsjdfk sjdkf lfke ljrl (contains dairy, soy)

    Parameters:
        filename (str): the name of the input file to open

    Returns:
        candidates (Dict[str, Set[str]]): a dict with keys that are the
            allergens (dairy, soy, etc.) and values that are the sets
            of all possible ingredients that could contain those
            allergens.
        counts (Counter): a count of how many lines each ingredient
            appears in.
    """

    candidates = dict()
    counts = Counter()
    with open(filename, "r") as f:
        for line in f:
            ingredients, _contains, allergens = line.partition(" (contains ")
            ingredients = set(ingredients.split())
            counts += Counter(ingredients)
            for allergen in allergens.strip(")\n").split(", "):
                if allergen in candidates:
                    candidates[allergen] &= ingredients
                else:
                    candidates[allergen] = ingredients.copy()

    return candidates, counts


def find_allergens(candidates):
    """Uses allergens that only have one possible ingredient that could
    contain them to rule that ingredient out for other allergens until
    all allergens have exactly one ingredient that could contain it.

    Returns a dict of allergen name to ingredient name to simplify things.
    """
    while not all(len(ingredients) == 1 for ingredients in candidates.values()):
        for allergen, ingredients in candidates.items():
            if len(ingredients) == 1:
                for other_al, other_ingredients in candidates.items():
                    if other_al != allergen:
                        other_ingredients -= ingredients
    return {allergen: ingredients.pop() for allergen, ingredients in candidates.items()}


def discard_unsafe_ingredient_counts(allergens, counts):
    """Given a set of decided allergens, removes each of these unsafe
    ingredients from a dict of counts so that only the counts of totally
    safe ingredients are retained.
    """
    for ingredient in allergens.values():
        del counts[ingredient]


def part1(counts):
    """Part 1 asks us to add up the number of occurrences of any totally
    safe ingredient.

    Assumes `counts` has already had its unsafe ingredients discarded.
    """
    return sum(counts.values())


def part2(allergens):
    """Part 2 wants the list of unsafe ingredients, sorted by their
    allergen alphabetically, joined together into one string by commas.
    """
    return ",".join(
        ingredient for (allergen, ingredient) in sorted(allergens.items())
    )

if __name__ == "__main__":
    test_candidates, test_counts = parse("data/day21_test.txt")
    test_allergens = find_allergens(test_candidates)
    discard_unsafe_ingredient_counts(test_allergens, test_counts)
    assert 5 == part1(test_counts)
    assert "mxmxvkd,sqjhc,fvjkl" == part2(test_allergens)
    print("All tests passed!")

    candidates, counts = parse("data/day21.txt")
    allergens = find_allergens(candidates)
    discard_unsafe_ingredient_counts(allergens, counts)
    print("Part 1:", part1(counts))
    print("Part 2:", part2(allergens))

use std::collections::{HashMap, HashSet};
use parser::*;

// -- model

#[derive(Debug, PartialEq)]
struct Food<'a> {
    ingredients: HashSet<&'a str>,
    allergens: HashSet<&'a str>
}

impl<'a> Food<'a> {
    fn new(ingredients: &[&'a str], allergens: &[&'a str]) -> Self {
        Food {
            ingredients: ingredients.iter().cloned().collect(),
            allergens: allergens.iter().cloned().collect()
        }
    }
}

struct Model<'a> {
    foods: Vec<Food<'a>>,
    ingredients_by_allergen: HashMap<&'a str, HashSet<&'a str>>,
}

impl<'a> Model<'a> {
    fn new(input: &'a str) -> Self {
        let ingredients = one_or_more(whitespace_wrap(word_ref));
        let allergens = word_ref
            .sep_by(match_literal(", "))
            .between(match_literal("(contains "), match_literal(")"));

        let food = pair(ingredients, allergens, |ingredients, allergens| Food {
            ingredients: ingredients.into_iter().collect(),
            allergens: allergens.into_iter().collect()
        });

        let foods = one_or_more(whitespace_wrap(food)).parse(input).unwrap().1;

        Model {
            foods,
            ingredients_by_allergen: HashMap::new()
        }
    }

    fn determine_allergens(&mut self) {
        self.associate_ingredients_with_allergens();
        while !self.is_fully_determined() {
            self.eliminate_duplicate_matches();
        }
    }

    fn associate_ingredients_with_allergens(&mut self) {
        for food in self.foods.iter() {
            for allergen in food.allergens.iter() {
                match self.ingredients_by_allergen.get_mut(allergen) {
                    Some(existing) => {
                        *existing = existing.intersection(&food.ingredients).cloned().collect();
                    }
                    None => {
                        self.ingredients_by_allergen.insert(allergen, food.ingredients.clone());
                    }
                }
            }
        }
    }

    fn is_fully_determined(&self) -> bool {
        self.ingredients_by_allergen.values().all(|ingredients| ingredients.len() < 2)
    }

    fn eliminate_duplicate_matches(&mut self) {
        let determined: HashSet<&'a str> = self.ingredients_by_allergen.values()
            .filter_map(|ingredients|
                if ingredients.len() == 1 { ingredients.iter().next() } else { None }
            ).cloned().collect();

        for ingredients in self.ingredients_by_allergen.values_mut().filter(|ings| ings.len() > 1) {
            *ingredients = ingredients.difference(&determined).cloned().collect();
        }
    }

    fn ingredients_with_allergen(&self) -> HashSet<&'a str> {
        self.ingredients_by_allergen.values().flat_map(|values| values.iter().cloned()).collect()
    }

    fn all_ingredients(&self) -> HashSet<&'a str> {
        self.foods.iter().flat_map(|food| food.ingredients.iter().cloned()).collect()
    }

    fn ingredients_with_no_allergen(&self) -> HashSet<&'a str> {
        self.all_ingredients()
            .difference(&self.ingredients_with_allergen())
            .cloned()
            .collect()
    }

    fn ingredients_alphabetically_by_allergen(&self) -> Vec<&'a str> {
        let mut allergens: Vec<&'a str> = self.ingredients_by_allergen.keys().cloned().collect();
        allergens.sort();
        allergens.iter().filter_map(|a| 
            self.ingredients_by_allergen.get(a).and_then(|ings| ings.iter().next())
        ).cloned().collect()
    }
}


// -- problems 

fn part1(model: &Model) -> usize {
    model.foods.iter().map(|food|
        food.ingredients.intersection(&model.ingredients_with_no_allergen()).count()
    ).sum()
}

fn part2(model: &Model) -> String {
    model.ingredients_alphabetically_by_allergen().join(",")
}

fn main() {
    let input = std::fs::read_to_string("./input.txt").unwrap();
    let mut model = Model::new(&input);
    model.determine_allergens();

    println!("part 1 {}", part1(&model));
    println!("part 2 {}", part2(&model));
}

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

    fn test_input() -> &'static str {
        "mxmxvkd kfcds sqjhc nhms (contains dairy, fish)
         trh fvjkl sbzzf mxmxvkd (contains dairy)
         sqjhc fvjkl (contains soy)
         sqjhc mxmxvkd sbzzf (contains fish)"        
    }

    #[test]
    fn test_parser() {
        let model = Model::new(test_input());
        assert_eq!(model.foods, vec![
            Food::new(
                &["mxmxvkd", "kfcds", "sqjhc", "nhms"],
                &["dairy", "fish"]
            ),
            Food::new(
                &["trh", "fvjkl", "sbzzf", "mxmxvkd"],
                &["dairy"]
            ),
            Food::new(
                &["sqjhc", "fvjkl"],
                &["soy"]
            ),
            Food::new(
                &["sqjhc", "mxmxvkd", "sbzzf"],
                &["fish"]
            )
        ]);
    }

    #[test]
    fn test_part1() {
        let mut model = Model::new(test_input());
        model.determine_allergens();
        assert_eq!(part1(&model), 5);
    }

    #[test]
    fn test_part2() {
        let mut model = Model::new(test_input());
        model.determine_allergens();
        assert_eq!(part2(&model), "mxmxvkd,sqjhc,fvjkl");
    }
}