Advent of Code 2019 - Day 6, orbits

aoc-2019-06.rs

use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

struct Tree<T> {
    head: Rc<RefCell<Node<T>>>,
}

impl Tree<String> {
    fn from_str(input: &str) -> Tree<String> {
        let tree = Tree {
            head: Rc::new(RefCell::new(Node::new("COM".to_string()))),
        };

        let mut map: HashMap<String, Rc<RefCell<Node<String>>>> = HashMap::new();
        map.insert("COM".into(), tree.head.clone());

        for line in input.lines() {
            let mut iter = line.split(')');
            let orbitee_name = iter.next().unwrap();

            let orbiter_name = iter.next().unwrap();
            let orbiter = Rc::new(RefCell::new(Node::new(orbiter_name.to_string())));

            map[orbitee_name]
                .borrow_mut()
                .children
                .push(orbiter.clone());
            map.insert(orbiter_name.into(), orbiter);
        }

        tree
    }
}

impl<T> Tree<T> {
    fn count_paths(&self) -> usize {
        self.head.borrow().count_paths(1)
    }
}

struct Node<T> {
    data: T,
    children: Vec<Rc<RefCell<Node<T>>>>,
}

impl<T> Node<T> {
    fn new(data: T) -> Node<T> {
        Node {
            data: data,
            children: Vec::new(),
        }
    }

    fn count_paths(&self, current: usize) -> usize {
        current * self.children.len()
            + self
                .children
                .iter()
                .map(|child| child.borrow().count_paths(current + 1))
                .sum::<usize>()
    }
}

fn main() {
    let input = "COM)B
B)C
C)D
D)E
E)F
B)G
G)H
D)I
E)J
J)K
K)L";

    assert_eq!(Tree::from_str(input).count_paths(), 42);
}