year2017/
day07.rs

1use std::collections::HashMap;
2use utils::prelude::*;
3use utils::str::TinyStr8;
4
5/// Finding the unbalanced subtree.
6#[derive(Clone, Debug)]
7pub struct Day07 {
8    programs: Vec<Program>,
9    bottom: usize,
10}
11
12#[derive(Clone, Debug)]
13struct Program {
14    name: TinyStr8,
15    weight: u32,
16    parent: Option<usize>,
17    children: Vec<usize>,
18}
19
20impl Day07 {
21    pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
22        let name = parser::tinystr8(u8::is_ascii_lowercase);
23        let lines = name
24            .with_suffix(" (")
25            .then(parser::u32().with_suffix(")"))
26            .then(name.repeat(", ", 0).with_prefix(" -> ".optional()))
27            .parse_lines(input)?;
28
29        let mut programs = lines
30            .iter()
31            .map(|&(name, weight, _)| Program {
32                name,
33                weight,
34                parent: None,
35                children: Vec::new(),
36            })
37            .collect::<Vec<_>>();
38
39        let name_map = lines
40            .iter()
41            .enumerate()
42            .map(|(index, &(name, _, _))| (name, index))
43            .collect::<HashMap<_, _>>();
44
45        for (parent, (_, _, children)) in lines.into_iter().enumerate() {
46            let children = children
47                .into_iter()
48                .map(|name| {
49                    if let Some(&child) = name_map.get(&name) {
50                        programs[child].parent = Some(parent);
51                        Ok(child)
52                    } else {
53                        Err(InputError::new(
54                            input,
55                            0,
56                            format!("program {name} missing on LHS"),
57                        ))
58                    }
59                })
60                .collect::<Result<Vec<_>, _>>()?;
61            programs[parent].children = children;
62        }
63
64        let Some(bottom) = programs.iter().position(|p| p.parent.is_none()) else {
65            return Err(InputError::new(
66                input,
67                0,
68                "expected one program to have no parent",
69            ));
70        };
71
72        Ok(Self { programs, bottom })
73    }
74
75    #[must_use]
76    pub fn part1(&self) -> String {
77        self.programs[self.bottom].name.to_string()
78    }
79
80    #[must_use]
81    pub fn part2(&self) -> u32 {
82        self.check(self.bottom).expect_err("tower is balanced")
83    }
84
85    fn check(&self, index: usize) -> Result<u32, u32> {
86        let program = &self.programs[index];
87        if program.children.is_empty() {
88            // Programs with no children are always balanced.
89            Ok(program.weight)
90        } else if program.children.len() < 3 {
91            // Programs with one child are balanced as there aren't multiple sub-towers to disagree.
92            // Programs with two children must also be balanced as it is impossible to tell which
93            // sub-tower is wrong if you only have two different values.
94            let first_weight = self.check(program.children[0])?;
95            let all_children = first_weight * program.children.len() as u32;
96            Ok(program.weight + all_children)
97        } else {
98            let first_weight = self.check(program.children[0])?;
99            let mut first_matches = 0;
100            let mut second_weight = None;
101            for &child in &program.children[1..] {
102                let weight = self.check(child)?;
103                if weight == first_weight {
104                    first_matches += 1;
105                } else if second_weight.is_none() {
106                    second_weight = Some((weight, child));
107                } else if second_weight.unwrap().0 != weight {
108                    panic!(
109                        "program {} has children with 3 different weights",
110                        program.name
111                    );
112                }
113            }
114
115            let Some((second_weight, second_index)) = second_weight else {
116                // All children match, this sub-tower is balanced
117                let all_children = first_weight * program.children.len() as u32;
118                return Ok(program.weight + all_children);
119            };
120
121            // Found the unbalanced sub-tower
122            let (correct_weight, wrong_weight, wrong_index) = if first_matches == 0 {
123                // First child wrong
124                (second_weight, first_weight, program.children[0])
125            } else {
126                // Second weight wrong
127                (first_weight, second_weight, second_index)
128            };
129
130            Err(correct_weight + self.programs[wrong_index].weight - wrong_weight)
131        }
132    }
133}
134
135examples!(Day07 -> (&'static str, u32) [
136    {file: "day07_example0.txt", part1: "tknk", part2: 60},
137]);
year2017/day07.rs

year2017/
day07.rs
