133 lines
3.7 KiB
Rust
133 lines
3.7 KiB
Rust
use std::{
|
|
cell::RefCell,
|
|
collections::{HashMap, HashSet},
|
|
io::{self, Read},
|
|
};
|
|
|
|
#[allow(unused_macros)]
|
|
macro_rules! todo {
|
|
($($tt:tt)*) => { unimplemented!($($tt)*) };
|
|
}
|
|
|
|
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
|
|
|
|
fn get_input_string() -> io::Result<String> {
|
|
let mut buffer = String::new();
|
|
io::stdin().read_to_string(&mut buffer)?;
|
|
Ok(buffer)
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
let orbit_list: Vec<Orbit> = get_input_string()?
|
|
.trim()
|
|
.lines()
|
|
.map(|line| {
|
|
let mut split = line.split(")");
|
|
Orbit {
|
|
parent: split.next().unwrap().to_string(),
|
|
child: split.next().unwrap().to_string(),
|
|
}
|
|
}).collect();
|
|
|
|
let mut orbit_graph = OrbitGraph::default();
|
|
for orbit in orbit_list.iter() {
|
|
orbit_graph.add_orbit(orbit.clone());
|
|
}
|
|
|
|
part1(&orbit_list, &orbit_graph);
|
|
part2(&orbit_list, &orbit_graph);
|
|
Ok(())
|
|
}
|
|
|
|
fn part1(orbit_list: &Vec<Orbit>, orbit_graph: &OrbitGraph) {
|
|
let bodies: HashSet<String> = orbit_list.iter()
|
|
.map(|o| o.child.clone())
|
|
.collect();
|
|
|
|
let total: usize = bodies.iter()
|
|
.map(|body| orbit_graph.indirect_count(body))
|
|
.sum();
|
|
println!("Part 1: {}", total);
|
|
}
|
|
|
|
fn part2(orbit_list: &Vec<Orbit>, orbit_graph: &OrbitGraph) {
|
|
let common = orbit_graph.common_parent("YOU", "SAN")
|
|
.expect("no common parent between YOU and SAN");
|
|
let common_dist = orbit_graph.indirect_count(common);
|
|
|
|
let hops1 = orbit_graph.indirect_count("YOU") - common_dist - 1;
|
|
let hops2 = orbit_graph.indirect_count("SAN") - common_dist - 1;
|
|
|
|
println!("Part 2: {}", hops1 + hops2);
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct Orbit {
|
|
parent: String,
|
|
child: String,
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
struct OrbitGraph {
|
|
graph: HashMap<String, Vec<String>>,
|
|
orbits: HashMap<String, String>,
|
|
count_cache: RefCell<HashMap<String, usize>>,
|
|
}
|
|
|
|
impl OrbitGraph {
|
|
fn add_orbit(&mut self, Orbit { parent, child }: Orbit) {
|
|
let child_list = self.graph.entry(parent.clone()).or_insert(Vec::new());
|
|
child_list.push(child.clone());
|
|
assert!(self.orbits.insert(child, parent).is_none());
|
|
}
|
|
|
|
fn indirect_count(&self, body: &str) -> usize {
|
|
{
|
|
let cache = self.count_cache.borrow();
|
|
if let Some(count) = cache.get(body) {
|
|
return *count;
|
|
}
|
|
}
|
|
|
|
if let Some(parent) = self.orbits.get(body) {
|
|
let count = 1 + self.indirect_count(parent);
|
|
{
|
|
let mut cache = self.count_cache.borrow_mut();
|
|
cache.insert(body.to_string(), count);
|
|
}
|
|
count
|
|
} else {
|
|
0
|
|
}
|
|
}
|
|
|
|
fn common_parent(&self, first: &str, second: &str) -> Option<&str> {
|
|
let first_parent = self.orbits.get(first);
|
|
let second_parent = self.orbits.get(second);
|
|
|
|
// One of the parents is the root, which is the common parent for everyone
|
|
if first_parent.is_none() || second_parent.is_none() {
|
|
return None;
|
|
}
|
|
|
|
let first_parent = first_parent.unwrap();
|
|
let second_parent= second_parent.unwrap();
|
|
|
|
// Common parent found
|
|
if first_parent == second_parent {
|
|
return Some(first_parent);
|
|
}
|
|
|
|
let depth1 = self.indirect_count(first);
|
|
let depth2 = self.indirect_count(second);
|
|
|
|
if depth1 < depth2 {
|
|
self.common_parent(first, second_parent)
|
|
} else if depth2 < depth1 {
|
|
self.common_parent(first_parent, second)
|
|
} else {
|
|
self.common_parent(first_parent, second_parent)
|
|
}
|
|
}
|
|
}
|