path: root/2019/src/bin/day_6.rs

use rpds::vector::Vector;
use std::fmt;
use std::io;
use std::io::prelude::*;
use std::iter::FromIterator;
use std::process;
use std::str::FromStr;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(name = "Day 6: Universal Orbit Map")]
/// Finds the minumum number of orbital transfers between two points.
///
/// Input is read from stdin, one direct orbit per line, in the format
/// `A)B` (B is orbiting A).
///
/// See https://adventofcode.com/2019/day/6 for details.
struct Opt {
    /// Debug checksum: Counts the total orbits
    #[structopt(short = "d", long = "debug")]
    debug: bool,
}

fn main() {
    let stdin = io::stdin();
    let opt = Opt::from_args();

    let orbits: OrbitalMap = stdin
        .lock()
        .lines()
        .map(|x| exit_on_failed_assertion(x, "Error reading input"))
        .map(|x| exit_on_failed_assertion(x.parse::<Orbit>(), "Input was not a valid orbit"))
        .collect();

    // eprintln!("{:#?}", orbits);

    if opt.debug {
        println!("{}", orbits.total_orbits());
    } else {
        println!("{}", orbits.orbital_transfers("YOU", "SAN"));
    }
}

fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
    match data {
        Ok(data) => data,
        Err(e) => {
            eprintln!("{}: {}", message, e);
            process::exit(1);
        }
    }
}

#[derive(Debug)]
struct StrError {
    str: String,
}

impl fmt::Display for StrError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.str)
    }
}
impl std::error::Error for StrError {}

#[derive(Debug, Clone)]
struct Orbit {
    a: String,
    b: String,
}

impl FromStr for Orbit {
    type Err = StrError;

    fn from_str(s: &str) -> Result<Self, StrError> {
        match s.split(')').collect::<Vec<_>>()[..] {
            [a, b] => Ok(Orbit {
                a: a.to_string(),
                b: b.to_string(),
            }),
            _ => Err(StrError {
                str: format!("{} is not a valid orbit description", s),
            }),
        }
    }
}

#[derive(Clone, Debug)]
struct OrbitalMap {
    id: String,
    depth: usize,
    orbiters: Vector<OrbitalMap>,
}

struct OrbitalMapBuilder {
    orbiters: Vector<OrbitalMap>,
    inserted_orbits: Vector<String>,
    pending_orbits: Vector<Orbit>,
}

impl FromIterator<Orbit> for OrbitalMap {
    fn from_iter<T: IntoIterator<Item = Orbit>>(iter: T) -> Self {
        iter.into_iter().collect::<OrbitalMapBuilder>().build()
    }
}

impl FromIterator<Orbit> for OrbitalMapBuilder {
    fn from_iter<T: IntoIterator<Item = Orbit>>(iter: T) -> Self {
        OrbitalMapBuilder {
            orbiters: Vector::new(),
            inserted_orbits: Vector::new(),
            pending_orbits: iter.into_iter().collect(),
        }
    }
}

impl OrbitalMapBuilder {
    fn build(self) -> OrbitalMap {
        if self.pending_orbits.is_empty() {
            OrbitalMap {
                id: ROOT.into(),
                depth: 0,
                orbiters: self.orbiters,
            }
        } else {
            self.pending_orbits
                .into_iter()
                .fold(
                    OrbitalMapBuilder {
                        pending_orbits: Vector::new(),
                        ..self
                    },
                    |acc, next| acc.insert(&next),
                )
                .build()
        }
    }

    fn insert(self, orbit: &Orbit) -> OrbitalMapBuilder {
        if orbit.a == ROOT {
            OrbitalMapBuilder {
                orbiters: self.orbiters.push_back(OrbitalMap::new(orbit.b.clone(), 1)),
                inserted_orbits: self.inserted_orbits.push_back(orbit.b.clone()),
                ..self
            }
        } else if self.inserted_orbits.iter().any(|o| *o == orbit.a) {
            OrbitalMapBuilder {
                orbiters: self
                    .orbiters
                    .into_iter()
                    .map(|map| OrbitalMap::insert(map.clone(), orbit))
                    .collect(),
                inserted_orbits: self.inserted_orbits.push_back(orbit.b.clone()),
                ..self
            }
        } else {
            OrbitalMapBuilder {
                pending_orbits: self.pending_orbits.push_back(orbit.clone()),
                ..self
            }
        }
    }
}

const ROOT: &str = "COM";

impl OrbitalMap {
    fn new(id: String, depth: usize) -> OrbitalMap {
        OrbitalMap {
            id,
            depth,
            orbiters: Vector::new(),
        }
    }

    fn insert(self, orbit: &Orbit) -> OrbitalMap {
        if orbit.a == self.id {
            if self.orbiters.iter().any(|o| o.id == orbit.b) {
                self
            } else {
                OrbitalMap {
                    orbiters: self
                        .orbiters
                        .push_back(OrbitalMap::new(orbit.b.clone(), self.depth + 1)),
                    ..self
                }
            }
        } else {
            OrbitalMap {
                orbiters: self
                    .orbiters
                    .into_iter()
                    .map(|map| OrbitalMap::insert(map.clone(), orbit))
                    .collect(),
                ..self
            }
        }
    }

    fn count_orbits(&self) -> usize {
        self.orbiters.len()
            + self
                .orbiters
                .iter()
                .map(|o| o.count_orbits())
                .sum::<usize>()
    }

    fn total_orbits(&self) -> usize {
        self.count_orbits()
            + self
                .orbiters
                .iter()
                .map(|o| o.total_orbits())
                .sum::<usize>()
    }

    fn find_depth(&self, id: &str) -> usize {
        if self.id == id {
            self.depth
        } else {
            // only the actual one will return non-zero from this
            self.orbiters.iter().map(|o| o.find_depth(id)).sum()
        }
    }

    fn contains(&self, id: &str) -> bool {
        self.id == id || self.orbiters.iter().any(|o| o.contains(id))
    }

    fn find_common_ancestor(&self, a: &str, b: &str) -> Option<OrbitalMap> {
        if !self.contains(a) || !self.contains(b) {
            None
        } else {
            self.orbiters
                .iter()
                .flat_map(|o| o.find_common_ancestor(a, b))
                .next()
                .or(Some(self.clone()))
        }
    }

    fn orbital_transfers(&self, from: &str, to: &str) -> usize {
        self.find_depth(from) + self.find_depth(to)
            - 2 * self
                .find_common_ancestor(from, to)
                .map(|o| o.depth)
                .unwrap_or(0)
            - 2
    }
}