82 lines
2 KiB
Rust
82 lines
2 KiB
Rust
#![feature(test)]
|
|
|
|
use regex::Regex;
|
|
use rustc_hash::FxHashMap;
|
|
|
|
struct Input {
|
|
bags: Vec<Vec<(u32, usize)>>,
|
|
start: usize,
|
|
}
|
|
|
|
fn setup(input: &str) -> Input {
|
|
let name_regex = Regex::new(r"^(.+?) bags").unwrap();
|
|
let content_regex = Regex::new(r"(\d+) (.+?) bags?").unwrap();
|
|
let names = input
|
|
.lines()
|
|
.enumerate()
|
|
.map(|(i, line)| {
|
|
let name = name_regex.captures(line).unwrap().get(1).unwrap().as_str();
|
|
(name, i)
|
|
})
|
|
.collect::<FxHashMap<_, _>>();
|
|
let bags = input
|
|
.lines()
|
|
.map(|line| {
|
|
content_regex
|
|
.captures_iter(line)
|
|
.map(|c| {
|
|
let count = c.get(1).unwrap().as_str().parse().unwrap();
|
|
let name = c.get(2).unwrap().as_str();
|
|
(count, names[name])
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
Input {
|
|
bags,
|
|
start: names["shiny gold"],
|
|
}
|
|
}
|
|
|
|
fn part1(input: &Input) -> u32 {
|
|
let mut rev = vec![vec![]; input.bags.len()];
|
|
for (i, bag) in input.bags.iter().enumerate() {
|
|
for &(_, j) in bag {
|
|
rev[j].push(i);
|
|
}
|
|
}
|
|
let mut out = 0;
|
|
let mut queue = vec![input.start];
|
|
let mut seen = vec![false; input.bags.len()];
|
|
seen[input.start] = true;
|
|
while let Some(p) = queue.pop() {
|
|
for &q in &rev[p] {
|
|
if !seen[q] {
|
|
seen[q] = true;
|
|
out += 1;
|
|
queue.push(q);
|
|
}
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
fn count_bags(input: &Input, counts: &mut [Option<u32>], bag: usize) -> u32 {
|
|
if let Some(out) = counts[bag] {
|
|
return out;
|
|
}
|
|
|
|
let mut out = 1;
|
|
for &(cnt, q) in &input.bags[bag] {
|
|
out += cnt * count_bags(input, counts, q);
|
|
}
|
|
counts[bag] = Some(out);
|
|
out
|
|
}
|
|
|
|
fn part2(input: &Input) -> u32 {
|
|
let mut counts = vec![None; input.bags.len()];
|
|
count_bags(input, &mut counts, input.start) - 1
|
|
}
|
|
|
|
aoc::main!(2020, 7, ex: 1, 2[b]);
|