2020_day07.rs

#![feature(or_insert_with_key, option_unwrap_none)]

use std::collections::HashMap;

use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::visit::{self, IntoNeighborsDirected, IntoNodeIdentifiers};
use petgraph::Direction;
use regex::Regex;

const INPUT: &str = include_str!("../input");

struct Bags(DiGraph<String, u32>);

impl Bags {
    fn count_incoming_shiny(&self) -> usize {
        let shiny_idx = self
            .0
            .node_indices()
            .find(|i| self.0[*i] == "shiny gold")
            .unwrap();

        self.count_incoming(shiny_idx)
    }
    fn count_containing_shiny(&self) -> usize {
        let shiny_idx = self
            .0
            .node_indices()
            .find(|i| self.0[*i] == "shiny gold")
            .unwrap();

        self.count_containing_bags(shiny_idx)
    }

    fn count_incoming(&self, idx: NodeIndex) -> usize {
        petgraph::visit::NodeFiltered::from_fn(&self.0, |node| {
            petgraph::algo::has_path_connecting(&self.0, node, idx, None)
        })
        .node_identifiers()
        .count()
            - 1
    }

    fn count_containing_bags(&self, idx: NodeIndex) -> usize {
        let filtered = petgraph::visit::NodeFiltered::from_fn(&self.0, |node| {
            petgraph::algo::has_path_connecting(&self.0, idx, node, None)
        });

        let mut weights: HashMap<&str, usize> = HashMap::new();
        let mut dfs = visit::DfsPostOrder::new(&filtered, idx);

        while let Some(next) = dfs.next(&filtered) {
            weights
                .insert(
                    &self.0[next],
                    filtered
                        .neighbors_directed(next, Direction::Outgoing)
                        .map(|neighbor| {
                            let edge = self.0.find_edge(next, neighbor).unwrap();

                            weights[self.0[neighbor].as_str()]
                                * *self.0.edge_weight(edge).unwrap() as usize
                        })
                        .sum::<usize>()
                        + 1,
                )
                .unwrap_none()
        }

        weights["shiny gold"] - 1
    }
}

fn parse(s: &str) -> Bags {
    let re = Regex::new(r"^(.+) bags contain (.+)\.$").unwrap();
    let inner_re = Regex::new(r"(?P<count>\d+) (?P<name>.+) bags?").unwrap();

    let mut graph = DiGraph::new();
    let mut map: HashMap<String, NodeIndex> = HashMap::new();

    for line in s.lines() {
        let captures = re.captures(line).unwrap();
        let bagname = captures[1].to_string();
        let other = &captures[2];

        let fidx = *map
            .entry(bagname)
            .or_insert_with_key(|k| graph.add_node(k.clone()));

        if other != "no other bags" {
            for inner_bag in other.split(", ") {
                let cap = inner_re.captures(inner_bag).unwrap();
                let bname = cap.name("name").unwrap().as_str().to_string();
                let count = cap.name("count").unwrap().as_str().parse::<u32>().unwrap();

                let sidx = *map
                    .entry(bname)
                    .or_insert_with_key(|k| graph.add_node(k.clone()));
                assert!(!graph.contains_edge(fidx, sidx));
                graph.add_edge(fidx, sidx, count);
            }
        }
    }

    Bags(graph)
}

fn main() {
    let bags = parse(INPUT);
    println!("A: {}", bags.count_incoming_shiny());
    println!("B: {}", bags.count_containing_shiny());
}

#[test]
fn test1() {
    let input = "light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.";

    let bags = parse(input);

    println!("{:?}", petgraph::dot::Dot::new(&bags.0));
    assert_eq!(4, bags.count_incoming_shiny());
}

#[test]
fn test2() {
    let input = "light red bags contain 1 bright white bag, 2 muted yellow bags.
dark orange bags contain 3 bright white bags, 4 muted yellow bags.
bright white bags contain 1 shiny gold bag.
muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
dark olive bags contain 3 faded blue bags, 4 dotted black bags.
vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
faded blue bags contain no other bags.
dotted black bags contain no other bags.";

    let bags = parse(input);

    println!("{:?}", petgraph::dot::Dot::new(&bags.0));
    assert_eq!(32, bags.count_containing_shiny());
}

#[test]
fn test3() {
    let input = "shiny gold bags contain 2 dark red bags.
dark red bags contain 2 dark orange bags.
dark orange bags contain 2 dark yellow bags.
dark yellow bags contain 2 dark green bags.
dark green bags contain 2 dark blue bags.
dark blue bags contain 2 dark violet bags.
dark violet bags contain no other bags.";

    let bags = parse(input);

    println!("{:?}", petgraph::dot::Dot::new(&bags.0));
    assert_eq!(126, bags.count_containing_shiny());
}