day7.rs

#[macro_use]
extern crate nom;

use nom::{alphanumeric, digit};
use std::str::FromStr;
use std::collections::{HashSet, HashMap, VecDeque};

// do_parse! lets you chain parsers together and give the intermediates a name.
// expr_res! lets you turn Result-returning functions into fake parsers.
// we use this to (step by step) turn a string into a usize.
named!(alphanumeric_usize(&str) -> usize,
    do_parse!(
        chars: digit >>
        weight: expr_res!(usize::from_str(chars)) >>
        (weight)
    )
);

// We return the alphanumeric chars as well as the weight.
// delimited! wraps the middle parse (we wrote that one!) and checks for
// but doesn't return the start and end (in this case, the brackets).
named!(name_and_weight(&str) -> (&str, usize),
    separated_pair!(
        alphanumeric,
        tag!(" "),
        delimited!(
            tag!("("),
            alphanumeric_usize,
            tag!(")")
        )
    )
);

// Handles the weird bit after the arrow.
// opt! makes this optional - if it fails we just don't include the tail bit.
// separated_nonempty_list_complete! kind of does what it says - it uses the
// first parser (tag!(",")) as a separator, and the second parser to parse
// the things in the list.
named!(carried(&str) -> Vec<&str>,
    do_parse!(
        tag!(" -> ") >>
        result: separated_nonempty_list_complete!(
            tag!(", "),
            alphanumeric
        ) >>
        (result)
    )
);

#[derive(PartialEq, Eq, Debug)]
struct Entry<'a> {
    name: &'a str,
    weight: usize,
    carried: Vec<&'a str>,
}

named!(entry(&str) -> Entry,
    do_parse!(
        name_and_weight: name_and_weight >>
        carried: opt!(carried) >>
        ({
            let (name, weight) = name_and_weight;
            let carried = carried.unwrap_or(Vec::new());
            (Entry {
                name, weight, carried
            })
        })
    )
);

named!(entries(&str) -> Vec<Entry>,
    separated_list_complete!(tag!("\n"), entry)
);

fn solution_1<'a>(entries: &[Entry<'a>]) -> &'a str {
    let mut carried_names: HashSet<&str> = HashSet::new();
    for carrier in entries.iter() {
        for carried in carrier.carried.iter() {
            carried_names.insert(carried);
        }
    }
    let uncarried_carrier = entries
        .iter()
        .filter(|carrier| !carried_names.contains(carrier.name))
        .nth(0)
        .unwrap();
    println!("root program: {:#?}", uncarried_carrier);
    uncarried_carrier.name
}

fn solution_2(entries: &[Entry]) {
    let mut name_to_program = HashMap::new();
    for program in entries.iter() {
        name_to_program.insert(program.name, program);
    }

    let mut name_to_weight = HashMap::new();

    // this is just a depth-first search done non-recursively.
    // the queue contains the current chain of programs that we're checking
    // the weight for.
    let mut work = VecDeque::new();
    let root = solution_1(entries);
    work.push_front(root);

    while let Some(name) = work.pop_front() {
        let program = name_to_program[name];
        if let Some(carried) = program.carried.get(0) {
            // either all our carried programs are ready, or none of them
            // are
            if name_to_weight.contains_key(carried) {
                let mut weight = program.weight;
                for carried in program.carried.iter() {
                    weight += name_to_weight[carried];
                }
                name_to_weight.insert(program.name, weight);
            } else {
                work.push_front(program.name);
                for carried in program.carried.iter() {
                    work.push_front(carried);
                }
            }
        } else {
            // we're a leaf node - just add ourselves
            name_to_weight.insert(program.name, program.weight);
        }
    }

    let mut parent_name = root;
    let mut target_weight = 0;
    loop {
        // need to find the problem child - the carried program with
        // a different weight, guaranteed to be unique.
        // we partition our carried programs by weight.
        let mut weight_to_carrieds = HashMap::new();
        let parent = name_to_program[parent_name];
        for carried in parent.carried.iter() {
            let weight = name_to_weight[carried];
            let carrieds = weight_to_carrieds.entry(weight).or_insert(Vec::new());
            carrieds.push(carried);
        }
        if let Some(bad_carried) =
            weight_to_carrieds
                .values()
                .filter(|v| v.len() == 1)
                .nth(0) {
            let bad_carried = bad_carried[0];
            // we have the name of the bad carried program - the one which
            // is alone in its' weight. The rest of the programs have the
            // same (correct) weight.
            let bad_weight = name_to_weight[bad_carried];
            parent_name = bad_carried;
            target_weight = weight_to_carrieds
                .keys()
                .cloned()
                .find(|weight| *weight != bad_weight)
                .unwrap();
            continue;
        } else {
            // all of our children are the same weight - that means we're
            // the problem. Just subtract the weight of our carried programs
            // from the target weight, and we have our answer!
            let carried_weight: usize = parent
                .carried
                .iter()
                .map(|carried| name_to_weight[carried])
                .sum();
            let target_weight = target_weight - carried_weight;
            println!("{} needs to be weight {}", parent.name, target_weight);
            break;
        }
    }
}

fn main() {
    let test = include_str!("../data/test.txt");
    let (_, test) = entries(test).unwrap();
    let test = test.as_ref();

    let real = include_str!("../data/real.txt");
    let (_, real) = entries(real).unwrap();
    let real = real.as_ref();

    solution_1(test);
    solution_1(real);

    solution_2(test);
    solution_2(real);
}

#[test]
fn check_parser() {
    use std::io::{BufRead, Cursor};
    use nom::IResult;
    let data = "hello (1234)";
    assert_eq!(name_and_weight(data), IResult::Done("", ("hello", 1234)));

    let data = " -> abc, de, f";
    assert_eq!(carried(data), IResult::Done("", vec!["abc", "de", "f"]));

    let data = "abc (2) -> def, e";
    assert_eq!(entry(data),
               IResult::Done("",
                             Entry {
                                 name: "abc",
                                 weight: 2,
                                 carried: vec!["def", "e"],
                             }));

    let data = "abc (2) -> def, e
def (345)
adff (4) -> arg, blarg";
    assert_eq!(entries(data),
               IResult::Done("",
                             vec![Entry {
                                      name: "abc",
                                      weight: 2,
                                      carried: vec!["def", "e"],
                                  },
                                  Entry {
                                      name: "def",
                                      weight: 345,
                                      carried: vec![],
                                  },
                                  Entry {
                                      name: "adff",
                                      weight: 4,
                                      carried: vec!["arg", "blarg"],
                                  }]));
}