Solution for Advent of Code 2017 Day 7

day07.rs

use std::time::Instant;

use rustc_hash::{FxHashMap, FxHashSet};

const INPUT: &str = include_str!("inputs/day07.txt");

#[derive(Clone)]
struct Program<'a> {
    weight: isize,
    total_weight: isize,
    children: Vec<&'a str>,
}

impl<'a> Program<'a> {
    fn parse(input: &'a str) -> (&'a str, Self) {
        let mut parts = input.split_whitespace();
        let name = parts.next().unwrap();
        let weight = parts
            .next()
            .unwrap()
            .trim_matches(['(', ')'])
            .parse::<isize>()
            .unwrap();

        let children = match parts.next() {
            Some(_) => parts.map(|p| p.trim_end_matches(',')).collect(),
            None => vec![],
        };
        (
            name,
            Program {
                weight,
                total_weight: 0,
                children,
            },
        )
    }
}

fn parse<'a>(input: &'a str) -> impl Iterator<Item = (&'a str, Program<'a>)> {
    input.lines().map(Program::parse)
}

fn p1(input: &str) -> &str {
    let (names, children) = parse(input).fold(
        (FxHashSet::default(), FxHashSet::default()),
        |(mut names, mut children), (name, program)| {
            names.insert(name);
            program.children.into_iter().for_each(|c| {
                children.insert(c);
            });
            (names, children)
        },
    );

    names.difference(&children).next().unwrap()
}

fn p2(input: &str) -> isize {
    let root = p1(input);
    let mut map = parse(input).collect::<FxHashMap<&str, Program>>();

    set_total_weights(&mut map, root, &mut FxHashMap::default());

    search(map, root, 0)
}

fn set_total_weights<'a>(
    map: &mut FxHashMap<&str, Program<'a>>,
    cur: &'a str,
    cache: &mut FxHashMap<&'a str, isize>,
) -> isize {
    if let Some(v) = cache.get(cur) {
        return *v;
    }
    let program = map[cur].clone();
    let children = program
        .children
        .iter()
        .map(|c| set_total_weights(map, c, cache))
        .sum::<isize>();
    let program = map.get_mut(cur).unwrap();
    program.total_weight = program.weight + children;
    cache.insert(cur, program.total_weight);
    program.total_weight
}

fn find_odd_one(children: &[isize]) -> Option<usize> {
    for (i, c) in children.iter().enumerate() {
        if children.iter().filter(|&m| m == c).count() == 1 {
            // sanity: odd one won't be a leaf
            return Some(i);
        }
    }
    None
}

fn search(map: FxHashMap<&str, Program>, cur: &str, other_weight: isize) -> isize {
    // Get all children weights.
    let program = map.get(&cur).unwrap().clone();
    let children = program
        .children
        .iter()
        .map(|c| map[c].total_weight)
        .collect::<Vec<_>>();

    // Find the odd-one-out of the children.
    if let Some(odd) = find_odd_one(&children) {
        let other = if odd == 0 { 1 } else { 0 };
        search(map, program.children[odd], children[other])
    } else {
        // If none, I'm the odd one.
        program.weight + (other_weight - program.total_weight)
    }
}

fn main() {
    let now = Instant::now();
    let solution = p1(INPUT);
    println!("p1 {:?} {}", now.elapsed(), solution);

    let now = Instant::now();
    let solution = p2(INPUT);
    println!("p2 {:?} {}", now.elapsed(), solution);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_p1() {
        let input = "pbga (66)\nxhth (57)\nebii (61)\nhavc (66)\nktlj (57)\nfwft (72) -> ktlj, cntj, xhth\nqoyq (66)\npadx (45) -> pbga, havc, qoyq\ntknk (41) -> ugml, padx, fwft\njptl (61)\nugml (68) -> gyxo, ebii, jptl\ngyxo (61)\ncntj (57)";
        assert_eq!(p1(input), "tknk");
    }
    #[test]
    fn test_p2() {
        let input = "pbga (66)\nxhth (57)\nebii (61)\nhavc (66)\nktlj (57)\nfwft (72) -> ktlj, cntj, xhth\nqoyq (66)\npadx (45) -> pbga, havc, qoyq\ntknk (41) -> ugml, padx, fwft\njptl (61)\nugml (68) -> gyxo, ebii, jptl\ngyxo (61)\ncntj (57)";
        assert_eq!(p2(input), 60);
    }
}