Clean up compute_dir_sizes

This commit is contained in:
pjht 2022-12-08 07:22:51 -06:00
parent 777b9dce30
commit 557b27a427

View File

@ -71,14 +71,12 @@ fn input_generator(input: &str) -> Vec<CliOperation> {
ops ops
} }
fn compute_dir_sizes(input: &[CliOperation]) -> Tree<(usize, usize)> { fn compute_dir_sizes(input: &[CliOperation]) -> Tree<usize> {
if input[0] != CliOperation::Cd("/".to_string()) { if input[0] != CliOperation::Cd("/".to_string()) {
panic!(); panic!();
} }
let mut tree = Tree::new(); let mut tree = Tree::new();
let mut current_node = tree let mut current_node = tree.insert(Node::new(0), InsertBehavior::AsRoot).unwrap();
.insert(Node::new((0, 0)), InsertBehavior::AsRoot)
.unwrap();
for op in input { for op in input {
match op { match op {
CliOperation::Cd(dir) => { CliOperation::Cd(dir) => {
@ -86,34 +84,30 @@ fn compute_dir_sizes(input: &[CliOperation]) -> Tree<(usize, usize)> {
current_node = tree.get(&current_node).unwrap().parent().unwrap().clone(); current_node = tree.get(&current_node).unwrap().parent().unwrap().clone();
} else { } else {
current_node = tree current_node = tree
.insert(Node::new((0, 0)), InsertBehavior::UnderNode(&current_node)) .insert(Node::new(0), InsertBehavior::UnderNode(&current_node))
.unwrap(); .unwrap();
}; };
} }
CliOperation::Ls(entries) => { CliOperation::Ls(entries) => {
*tree.get_mut(&current_node).unwrap().data_mut() = ( *tree.get_mut(&current_node).unwrap().data_mut() = entries
entries
.iter() .iter()
.filter_map(DirectoryEntry::as_file) .filter_map(DirectoryEntry::as_file)
.map(|file| file.size) .map(|file| file.size)
.sum::<usize>(), .sum::<usize>();
0,
)
} }
} }
} }
let post_order_ids = tree for id in tree
.traverse_post_order_ids(tree.root_node_id().unwrap()) .traverse_post_order_ids(tree.root_node_id().unwrap())
.unwrap() .unwrap()
.collect::<Vec<_>>(); {
for id in post_order_ids {
let child_sum = tree let child_sum = tree
.children(&id) .children(&id)
.unwrap() .unwrap()
.map(|node| node.data().1) .map(|node| node.data())
.sum::<usize>(); .sum::<usize>();
let self_sum = tree.get(&id).unwrap().data().0; let self_sum = tree.get(&id).unwrap().data();
tree.get_mut(&id).unwrap().data_mut().1 = child_sum + self_sum; *tree.get_mut(&id).unwrap().data_mut() = child_sum + self_sum;
} }
tree tree
} }
@ -123,7 +117,7 @@ fn solve_part1(input: &[CliOperation]) -> usize {
let tree = compute_dir_sizes(input); let tree = compute_dir_sizes(input);
tree.traverse_post_order(tree.root_node_id().unwrap()) tree.traverse_post_order(tree.root_node_id().unwrap())
.unwrap() .unwrap()
.map(|node| node.data().1) .map(|node| *node.data())
.filter(|&size| size <= 100_000) .filter(|&size| size <= 100_000)
.sum() .sum()
} }
@ -131,11 +125,11 @@ fn solve_part1(input: &[CliOperation]) -> usize {
#[aoc(day7, part2)] #[aoc(day7, part2)]
fn solve_part2(input: &[CliOperation]) -> usize { fn solve_part2(input: &[CliOperation]) -> usize {
let tree = compute_dir_sizes(input); let tree = compute_dir_sizes(input);
let free_space = 70_000_000 - tree.get(tree.root_node_id().unwrap()).unwrap().data().1; let free_space = 70_000_000 - tree.get(tree.root_node_id().unwrap()).unwrap().data();
let necessary_to_free = 30_000_000 - free_space; let necessary_to_free = 30_000_000 - free_space;
tree.traverse_post_order(tree.root_node_id().unwrap()) tree.traverse_post_order(tree.root_node_id().unwrap())
.unwrap() .unwrap()
.map(|node| node.data().1) .map(|node| *node.data())
.filter(|&size| size >= necessary_to_free) .filter(|&size| size >= necessary_to_free)
.min() .min()
.unwrap() .unwrap()