diff --git a/src/main.rs b/src/main.rs index eb062b8..6d037e8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,6 @@ #![feature(test)] + +use std::collections::HashMap; extern crate test; const INPUTS: [&str; 2] = [ @@ -12,6 +14,19 @@ fn parse(input: &'static str) -> impl Iterator { input } +#[derive(Debug)] +enum Node { + File(u32), + Directory(Directory), +} + +#[derive(Debug)] +struct Directory { + children: HashMap, + parent: *mut Node, + size: u32, +} + fn main() { for input in INPUTS.iter() { let output = parse(input); @@ -22,56 +37,111 @@ fn main() { } fn solution(input: impl Iterator) -> u32 { - let mut out = Vec::with_capacity(100); - let mut stack = Vec::with_capacity(100); + let mut tree = Node::Directory(Directory { + children: HashMap::new(), + parent: std::ptr::null_mut(), + size: 0, + }); - let mut current_folder_size = 0; + let mut current = &mut tree; for line in input { match &line[..4] { - "$ ls" | "dir " => continue, + "$ ls" => continue, + + "dir " => { + let dir = &line[4..]; + let cptr = current as *mut Node; + + let Node::Directory(Directory { children, ..}) = current else { unreachable!()}; + + children.insert( + dir.to_string(), + Node::Directory(Directory { + children: HashMap::new(), + parent: cptr, + size: 0, + }), + ); + } "$ cd" => match &line[5..6] { // we are supposed to match on .. but this is fine "." => { - let v = stack.pop().unwrap(); - out.push(current_folder_size); - current_folder_size += v; + let Node::Directory(Directory { parent , .. }) = *current else { + unreachable!("current is not a directory"); + }; + + current = unsafe { &mut *parent }; } - "/" => continue, + "/" => current = &mut tree, _ => { - stack.push(current_folder_size); - current_folder_size = 0; + let dir = &line[5..]; + let Node::Directory(Directory { children, ..}) = current else { + unreachable!("current is not a directory"); + }; + + current = children.get_mut(dir).unwrap(); } }, _ => { - let size = line + let (size, name) = line.split_once(' ').unwrap(); + let fsize = size .bytes() .take_while(|&c| c != b' ') .fold(0u32, |a, x| a * 10 + (x - b'0') as u32); - current_folder_size += size; + + if let Node::Directory(v) = current { + v.children + .entry(name.to_string()) + .or_insert(Node::File(fsize)); + } } } } - while let Some(v) = stack.pop() { - out.push(current_folder_size); - current_folder_size += v; + compute_dir_size(&mut tree); + + part1(&tree) +} + +fn part1(node: &Node) -> u32 { + match node { + &Node::File(_) => 0, + + Node::Directory(dir) => { + let mut sum = 0; + if dir.size <= 100000 { + sum += dir.size; + } + + for v in dir.children.values() { + sum += part1(v); + } + + sum + } } +} - out.push(current_folder_size); +fn compute_dir_size(node: &mut Node) -> u32 { + match node { + &mut Node::File(v) => v, - let available_space = 70000000; - let required_space = 30000000; - let used_space = *out.last().unwrap(); + Node::Directory(dir) => { + let mut sum = 0; - out.into_iter() - .filter(|&c| c >= (required_space + used_space - available_space)) - .min() - .unwrap() + for v in dir.children.values_mut() { + sum += compute_dir_size(v); + } + + dir.size = sum; + sum + } + } } #[bench]