1
0
Fork 0
aoc2020/src/bin/day7.rs

124 lines
3.3 KiB
Rust

use std::collections::{HashMap, HashSet, VecDeque};
const INPUT: &str = include_str!("../../data/day7");
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
struct BagId(usize);
#[derive(Clone, Debug)]
struct Bag {
name: String,
contained_by: HashMap<BagId, usize>,
contains: HashMap<BagId, usize>,
contained_bags: Option<u64>,
}
impl Bag {
fn new(bag: &str) -> Self {
Self {
name: bag.into(),
contained_by: HashMap::new(),
contains: HashMap::new(),
contained_bags: None,
}
}
}
#[derive(Clone, Debug)]
struct Bags {
name_to_id: HashMap<String, BagId>,
bags: Vec<Bag>,
}
impl Bags {
fn get_bag_id(&mut self, bag: &str) -> BagId {
let bags = &mut self.bags; // for access in closure
*self.name_to_id.entry(bag.to_string()).or_insert_with(|| {
let id = BagId(bags.len());
bags.push(Bag::new(bag));
id
})
}
fn add_bag_spec(&mut self, bag: &str, contains: &[(&str, usize)]) {
let bag_id = self.get_bag_id(bag);
let contains: HashMap<BagId, usize> = contains.into_iter().map(|(name, amount)| {
(self.get_bag_id(name), *amount)
}).collect();
for (&id, &amount) in &contains {
assert!(amount > 0);
self.bags[id.0].contained_by.insert(bag_id, amount);
}
let bag = &mut self.bags[bag_id.0];
assert!(bag.contains.is_empty());
bag.contains = contains;
}
fn add_bag_spec_str(&mut self, line: &str) {
use std::convert::TryInto;
let line = line.trim().strip_suffix(".").expect("need trailing dot");
let [bag_name, specs]: [&str; 2] = line.splitn(2, " bags contain ").collect::<Vec<_>>().try_into().expect("invalid line");
let specs = if specs != "no other bags" {
specs.split(", ").map(|spec| {
if let Some(spec) = spec.strip_suffix(" bag") {
let spec = spec.strip_prefix("1 ").expect("singular bag requires amount 1");
(spec, 1)
} else {
let spec = spec.strip_suffix(" bags").expect("either singular or plural required");
let [amount, spec]: [&str; 2] = spec.splitn(2, " ").collect::<Vec<_>>().try_into().expect("invalid spec");
(spec, amount.parse::<usize>().unwrap())
}
}).collect::<Vec<_>>()
} else {
Vec::new()
};
self.add_bag_spec(bag_name, &specs);
}
fn parse(data: &str) -> Self {
let mut this = Self {
name_to_id: HashMap::new(),
bags: Vec::new(),
};
for line in data.lines() {
this.add_bag_spec_str(line)
}
this
}
fn find_bag(&self, name: &str) -> Option<BagId> {
self.name_to_id.get(name).cloned()
}
fn contained_bags(&mut self, bag: BagId) -> u64 {
if let Some(n) = self.bags[bag.0].contained_bags {
return n;
}
let contains = self.bags[bag.0].contains.clone();
let n = contains.into_iter().map(|(inner_bag, amount)| {
(amount as u64) * (self.contained_bags(inner_bag) + 1)
}).sum::<u64>();
self.bags[bag.0].contained_bags = Some(n);
n
}
}
fn main() {
let mut bags = Bags::parse(INPUT);
let shiny_gold = bags.find_bag("shiny gold").unwrap();
let mut walk = VecDeque::new();
let mut contain_shiny_gold = HashSet::new();
walk.push_back(shiny_gold);
while let Some(bag_id) = walk.pop_front() {
if contain_shiny_gold.insert(bag_id) {
for contains in bags.bags[bag_id.0].contained_by.keys() {
walk.push_back(*contains);
}
}
}
println!("Bags containing shiny gold: {}", contain_shiny_gold.len() - 1);
println!("shiny gold contains: {}", bags.contained_bags(shiny_gold));
}