year2017/
day07.rs

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

year2017/
day07.rs
