advent_of_code/year2017/
day07.rs

1use crate::input::Input;
2use std::collections::{HashMap, HashSet};
3
4type ProgramId = usize;
5
6#[derive(Clone)]
7struct Program<'a> {
8    name: &'a str,
9    weight: u32,
10    children: Vec<ProgramId>,
11}
12
13struct ProgramTree<'a> {
14    nodes: Vec<Program<'a>>,
15    root_node: ProgramId,
16}
17
18impl ProgramTree<'_> {
19    fn total_weight(&self, program_id: ProgramId) -> u32 {
20        let program = &self.nodes[program_id];
21        program.weight
22            + program
23                .children
24                .iter()
25                .map(|&child| self.total_weight(child))
26                .sum::<u32>()
27    }
28}
29
30impl<'a> ProgramTree<'a> {
31    fn parse(input_string: &'a str) -> Result<Self, String> {
32        let mut nodes = Vec::new();
33        let mut name_to_node: HashMap<&str, ProgramId> = HashMap::new();
34
35        for (line_index, line) in input_string.lines().enumerate() {
36            let parts: Vec<&str> = line.split(" -> ").collect();
37            let (name, weight_str) = parts[0].split_once(' ').ok_or_else(|| {
38                format!(
39                    "Line {}: Invalid format, expected '$NAME ($WEIGHT)'",
40                    line_index + 1
41                )
42            })?;
43            let weight = weight_str
44                .replace(['(', ')'], "")
45                .parse::<u32>()
46                .map_err(|error| format!("Line {}: Invalid weight ({})", line_index + 1, error))?;
47
48            let program_id = nodes.len();
49            let program = Program {
50                name,
51                weight,
52                children: Vec::new(),
53            };
54            nodes.push(program);
55            name_to_node.insert(name, program_id);
56        }
57
58        for line in input_string.lines() {
59            let parts: Vec<&str> = line.split(" -> ").collect();
60            let name_weight_parts: Vec<&str> = parts[0].split(' ').collect();
61            if parts.len() == 2 {
62                let name = name_weight_parts[0];
63                let parent = *name_to_node
64                    .get_mut(name)
65                    .ok_or("Internal error - no node for name")?;
66                for child_name in parts[1].trim().split(", ") {
67                    let child_ref = *name_to_node
68                        .get(child_name)
69                        .ok_or("Invalid input - no node for name")?;
70                    nodes[parent].children.push(child_ref);
71                }
72            }
73        }
74
75        let all_program_ids: HashSet<&ProgramId> = name_to_node.values().collect();
76        let children: HashSet<&ProgramId> = name_to_node
77            .values()
78            .flat_map(|&child_program_id| nodes[child_program_id].children.iter())
79            .collect();
80
81        let roots: Vec<&&ProgramId> = all_program_ids.difference(&children).collect();
82        let root_node = **roots[0];
83        if roots.len() == 1 {
84            Ok(Self { nodes, root_node })
85        } else {
86            Err("No single root found".to_string())
87        }
88    }
89}
90
91pub fn solve(input: &Input) -> Result<String, String> {
92    let tree = ProgramTree::parse(input.text)?;
93    if input.is_part_one() {
94        Ok(tree.nodes[tree.root_node].name.to_string())
95    } else {
96        fixup_weight(tree.root_node, &tree)
97            .map(|value| value.to_string())
98            .ok_or_else(|| "No solution found".to_string())
99    }
100}
101
102fn fixup_weight(program_id: ProgramId, tree: &ProgramTree) -> Option<u32> {
103    let program = &tree.nodes[program_id];
104    if program.children.len() > 1 {
105        if let (Some(lone_weight), desired_weight) = program
106            .children
107            .iter()
108            .fold(HashMap::new(), |mut acc, &child_id| {
109                *acc.entry(tree.total_weight(child_id)).or_insert(0) += 1;
110                acc
111            })
112            .iter()
113            .fold((None, 0), |acc, (&weight, &occurrences)| {
114                if occurrences == 1 {
115                    (Some(weight), acc.1)
116                } else {
117                    (acc.0, weight)
118                }
119            })
120        {
121            if let Some(&child_id) = program
122                .children
123                .iter()
124                .find(|&&p| tree.total_weight(p) == lone_weight)
125            {
126                return fixup_weight(child_id, tree).or_else(|| {
127                    let total_weight = tree.total_weight(child_id);
128                    let child = &tree.nodes[child_id];
129                    Some(desired_weight - (total_weight - child.weight))
130                });
131            }
132        }
133    } else if !program.children.is_empty() {
134        if let Some(value) = fixup_weight(program.children[0], tree) {
135            return Some(value);
136        }
137    }
138    None
139}
140
141#[test]
142fn tests() {
143    use crate::input::{test_part_one, test_part_two};
144    let real_input = include_str!("day07_input.txt");
145    test_part_one!(real_input => "veboyvy".to_string());
146    test_part_two!(real_input => "749".to_string());
147}
advent_of_code/year2017/day07.rs

advent_of_code/year2017/
day07.rs
