From 08ddc927b812d5f38666b3dd7de843e942710b13 Mon Sep 17 00:00:00 2001 From: Jesse Luehrs Date: Thu, 8 Dec 2022 01:21:54 -0500 Subject: simplify --- src/2022/7/mod.rs | 117 ++++++++++----------------------- src/main.rs | 1 + src/prelude.rs | 1 + src/tree.rs | 189 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 226 insertions(+), 82 deletions(-) create mode 100644 src/tree.rs diff --git a/src/2022/7/mod.rs b/src/2022/7/mod.rs index bdc5192..9fbaf97 100644 --- a/src/2022/7/mod.rs +++ b/src/2022/7/mod.rs @@ -3,74 +3,35 @@ use crate::prelude::*; -enum Node { - Dir(HashMap), +pub enum Node { + Dir, File(i64), } impl Node { fn new_dir() -> Self { - Self::Dir(HashMap::new()) + Self::Dir } - + fn new_file(size: i64) -> Self { Self::File(size) } - + fn size(&self) -> i64 { match self { + Self::Dir => 0, Self::File(size) => *size, - Self::Dir(children) => children - .values() - .map(|child| child.size()) - .sum(), - } - } - - fn iter(&self) -> NodeIter<'_> { - let mut next = std::collections::VecDeque::new(); - next.push_back(("", self)); - NodeIter { next } - } - - fn resolve_path(&mut self, path: &[String]) -> &mut Self { - let mut pwd = self; - for name in path { - let Node::Dir(children) = pwd - else { panic!("invalid path") }; - - pwd = children.get_mut(name).unwrap(); - } - pwd - } -} - -struct NodeIter<'a> { - next: std::collections::VecDeque<(&'a str, &'a Node)>, -} - -impl<'a> Iterator for NodeIter<'a> { - type Item = (&'a str, &'a Node); - - fn next(&mut self) -> Option { - let Some(next) = self.next.pop_back() - else { return None }; - if let Node::Dir(children) = next.1 { - for (name, child) in children.iter() { - self.next.push_front((name, child)); - } } - Some(next) } } -pub struct DirTree { - root: Node, +fn tree_size(tree: &Tree) -> i64 { + tree.bfs().map(|(_, tree)| tree.data().size()).sum() } -pub fn parse(fh: File) -> Result { +pub fn parse(fh: File) -> Result> { let mut lines = parse::lines(fh).peekable(); - let mut root = Node::new_dir(); + let mut root = Tree::new(Node::Dir); let mut path = vec![]; while let Some(cmdline) = lines.by_ref().next() { @@ -82,7 +43,7 @@ pub fn parse(fh: File) -> Result { let Some(cmd) = captures.get(1) else { bail!("no cmd found in cmdline: '{}'", cmdline) }; - let pwd = root.resolve_path(&path); + let pwd = root.at_mut(&path).unwrap(); match cmd.as_str() { "cd" => { @@ -94,11 +55,8 @@ pub fn parse(fh: File) -> Result { } "/" => path = vec![], dirname => { - let Node::Dir(children) = pwd - else { panic!("invalid path") }; - - if !children.contains_key(dirname) { - children.insert( + if pwd.at(&[dirname.to_string()]).is_none() { + pwd.add_child( dirname.to_string(), Node::new_dir(), ); @@ -107,45 +65,40 @@ pub fn parse(fh: File) -> Result { } } } - "ls" => { - loop { - let Node::Dir(children) = pwd - else { panic!("invalid path") }; - - let Some(lsline) = lines.peek() + "ls" => loop { + let Some(lsline) = lines.peek() else { break }; - let Some(captures) = ®ex_captures!( + let Some(captures) = ®ex_captures!( r"^(dir|[0-9]+) (.*)$", lsline ) else { break }; - let Some(data) = captures.get(1) + let Some(data) = captures.get(1) else { break }; - let Some(name) = captures.get(2) + let Some(name) = captures.get(2) else { break }; - let node = match data.as_str() { - "dir" => Node::new_dir(), - size => Node::new_file(size.parse().unwrap()), - }; - children.insert(name.as_str().to_string(), node); - lines.next(); - } - } + let node = match data.as_str() { + "dir" => Node::new_dir(), + size => Node::new_file(size.parse().unwrap()), + }; + pwd.add_child(name.as_str().to_string(), node); + lines.next(); + }, _ => unreachable!(), } } - Ok(DirTree { root }) + Ok(root) } -pub fn part1(tree: DirTree) -> Result { +pub fn part1(tree: Tree) -> Result { let mut total = 0; - for (name, node) in tree.root.iter() { - if matches!(node, Node::File(_)) { + for (_, tree) in tree.bfs() { + if matches!(tree.data(), Node::File(_)) { continue; } - let size = node.size(); + let size = tree_size(tree); if size <= 100_000 { total += size; } @@ -153,16 +106,16 @@ pub fn part1(tree: DirTree) -> Result { Ok(total) } -pub fn part2(tree: DirTree) -> Result { - let total = tree.root.size(); +pub fn part2(tree: Tree) -> Result { + let total = tree_size(&tree); let free = 70_000_000 - total; let needed = 30_000_000 - free; let mut possible = vec![]; - for (name, node) in tree.root.iter() { - if matches!(node, Node::File(_)) { + for (_, tree) in tree.bfs() { + if matches!(tree.data(), Node::File(_)) { continue; } - let size = node.size(); + let size = tree_size(tree); if size > needed { possible.push(size); } diff --git a/src/main.rs b/src/main.rs index e04fedb..490fa02 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ pub mod graph; pub mod grid; pub mod parse; pub mod prelude; +pub mod tree; #[path = "2020/mod.rs"] mod year2020; diff --git a/src/prelude.rs b/src/prelude.rs index 6c53874..046b4fb 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -1,6 +1,7 @@ pub use crate::graph::Graph as _; pub use crate::grid::{Col, Grid, Row}; pub use crate::parse; +pub use crate::tree::Tree; pub use std::cmp::Ordering; pub use std::collections::VecDeque; diff --git a/src/tree.rs b/src/tree.rs new file mode 100644 index 0000000..3c727dd --- /dev/null +++ b/src/tree.rs @@ -0,0 +1,189 @@ +use crate::prelude::*; + +#[derive(Debug)] +pub struct Tree { + data: V, + children: HashMap>, +} + +impl Tree +where + K: std::hash::Hash + std::cmp::Eq + std::clone::Clone, +{ + pub fn new(data: V) -> Self { + Self { + data, + children: HashMap::new(), + } + } + + pub fn data(&self) -> &V { + &self.data + } + + pub fn data_mut(&mut self) -> &mut V { + &mut self.data + } + + pub fn add_child(&mut self, name: K, data: V) { + self.children.insert(name, Self::new(data)); + } + + pub fn at<'a, I>(&'a self, path: I) -> Option<&'a Self> + where + I: IntoIterator, + { + let mut ret = self; + for entry in path { + let Some(child) = ret.children.get(entry) + else { return None }; + ret = child; + } + Some(ret) + } + + pub fn at_mut<'a, I>(&'a mut self, path: I) -> Option<&'a mut Self> + where + I: IntoIterator, + { + let mut ret = self; + for entry in path { + let Some(child) = ret.children.get_mut(entry) + else { return None }; + ret = child; + } + Some(ret) + } + + pub fn bfs(&self) -> BFS<'_, K, V> { + let mut next = std::collections::VecDeque::new(); + next.push_front((vec![], self)); + BFS { next } + } + + pub fn dfs(&self) -> DFS<'_, K, V> { + DFS { + next: vec![(vec![], self)], + } + } +} + +pub struct BFS<'a, K, V> +where + K: std::cmp::Eq + std::hash::Hash + std::clone::Clone, +{ + next: std::collections::VecDeque<(Vec, &'a Tree)>, +} + +impl<'a, K, V> Iterator for BFS<'a, K, V> +where + K: std::cmp::Eq + std::hash::Hash + std::clone::Clone, +{ + type Item = (Vec, &'a Tree); + + fn next(&mut self) -> Option { + let Some(next) = self.next.pop_front() + else { return None }; + + self.next.extend(next.1.children.iter().map(|(k, v)| { + ( + next.0 + .clone() + .into_iter() + .chain(std::iter::once(k.clone())) + .collect(), + v, + ) + })); + Some(next) + } +} + +pub struct DFS<'a, K, V> +where + K: std::cmp::Eq + std::hash::Hash + std::clone::Clone, +{ + next: Vec<(Vec, &'a Tree)>, +} + +impl<'a, K, V> Iterator for DFS<'a, K, V> +where + K: std::cmp::Eq + std::hash::Hash + std::clone::Clone, +{ + type Item = (Vec, &'a Tree); + + fn next(&mut self) -> Option { + let Some(next) = self.next.pop() + else { return None }; + + self.next.extend(next.1.children.iter().map(|(k, v)| { + ( + next.0 + .clone() + .into_iter() + .chain(std::iter::once(k.clone())) + .collect(), + v, + ) + })); + Some(next) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic() { + enum Entry { + Node, + Leaf(u32), + } + + impl Entry { + fn value(&self) -> u32 { + match self { + Self::Node => 0, + Self::Leaf(value) => *value, + } + } + } + + let mut tree = Tree::new(Entry::Node); + tree.add_child("foo", Entry::Node); + tree.add_child("bar", Entry::Leaf(23)); + tree.at_mut(&["foo"]) + .unwrap() + .add_child("baz", Entry::Leaf(42)); + + let bfs = tree.bfs().map(|(path, _)| path).collect::>(); + assert_eq!(bfs[0], Vec::<&str>::new()); + assert!(bfs[1] == vec!["foo"] || bfs[1] == vec!["bar"]); + assert!(bfs[2] == vec!["foo"] || bfs[2] == vec!["bar"]); + assert!(bfs[1] != bfs[2]); + assert_eq!(bfs[3], vec!["foo", "baz"]); + + let dfs = tree.dfs().map(|(path, _)| path).collect::>(); + assert_eq!(dfs[0], Vec::<&str>::new()); + assert!(dfs[1] == vec!["foo"] || dfs[1] == vec!["bar"]); + if dfs[1] == vec!["foo"] { + assert_eq!(dfs[2], vec!["foo", "baz"]); + assert_eq!(dfs[3], vec!["bar"]); + } else if dfs[1] == vec!["bar"] { + assert_eq!(dfs[2], vec!["foo"]); + assert_eq!(dfs[3], vec!["foo", "baz"]); + } else { + unreachable!() + } + + assert_eq!( + tree.bfs().map(|(_, tree)| tree.data().value()).sum::(), + 65 + ); + assert_eq!( + tree.dfs().map(|(_, tree)| tree.data().value()).sum::(), + 65 + ); + } +} -- cgit v1.2.3-54-g00ecf