aoc2019/day6/src/main.rs

87 lines
2.4 KiB
Rust

use std::cell::Cell;
use std::collections::{HashMap, HashSet};
fn parse(s: &str) -> impl Iterator<Item=(&str, &str)> {
s.split_ascii_whitespace().map(|l| {
let mut x = l.splitn(2, ')');
let a = x.next().unwrap();
let b = x.next().unwrap();
(a, b)
})
}
fn count_orbits_get(orbit_relations: &HashMap<&str, (&str, Cell<Option<u32>>)>, object: &str) -> u32 {
if object == "COM" { return 0; }
let around = &orbit_relations[object];
if let Some(n) = around.1.get() { return n; }
let n = count_orbits_get(orbit_relations, around.0) + 1;
around.1.set(Some(n));
n
}
fn count_orbits(input: &str) -> u32 {
let orbit_relations: HashMap<&str, (&str, Cell<Option<u32>>)> = parse(input).map(|(a, b)| (b, (a, Cell::new(None)))).collect();
let mut sum = 0;
for (_object, around) in &orbit_relations {
if let Some(n) = around.1.get() {
sum += n;
} else {
let n = count_orbits_get(&orbit_relations, around.0) + 1;
around.1.set(Some(n));
sum += n;
}
}
// println!("Map: {:?}", orbit_relations);
sum
}
fn count_orbit_transfers(input: &str) -> u32 {
let orbit_relations: HashMap<&str, (&str, Cell<Option<u32>>)> = parse(input).map(|(a, b)| (b, (a, Cell::new(None)))).collect();
for (_object, around) in &orbit_relations {
if around.1.get().is_none() {
let n = count_orbits_get(&orbit_relations, around.0) + 1;
around.1.set(Some(n));
}
}
// println!("Map: {:?}", orbit_relations);
let mut you = HashSet::new();
let mut object = "YOU";
loop {
if "COM" == object { break; }
object = orbit_relations[object].0;
// println!("YOU orbit (...) {}", object);
you.insert(object);
}
object = "SAN";
loop {
if you.contains(&object) {
let you = orbit_relations["YOU"].1.get().unwrap();
let san = orbit_relations["SAN"].1.get().unwrap();
let bottom = orbit_relations[object].1.get().unwrap();
return (you - bottom - 1) + (san - bottom - 1);
}
object = orbit_relations[object].0;
// println!("SAN orbit (...) {}", object);
}
}
fn main() {
let input = include_str!("input.txt");
println!("Orbit sum: {}", count_orbits(input));
println!("Orbit transfers: {}", count_orbit_transfers(input));
}
#[cfg(test)]
mod day6test {
use super::*;
#[test]
fn examples1() {
assert_eq!(count_orbits("COM)B B)C C)D D)E E)F B)G G)H D)I E)J J)K K)L"), 42);
}
#[test]
fn examples2() {
assert_eq!(count_orbit_transfers("COM)B B)C C)D D)E E)F B)G G)H D)I E)J J)K K)L K)YOU I)SAN"), 4);
}
}