feat: add tree pruning for budget challenge

Author: Curstantine

src/challenge/challenge_07.rs

@@ -0,0 +1,141 @@
+//! Challenge 07: Tree Pruning for Budget
+use anyhow::Result;
+use colored::Colorize;
+
+use crate::utils::split_whitespace_and_parse;
+
+#[derive(Clone)]
+struct Node {
+    value: u32,
+    cost: u32,
+    children: Vec<usize>,
+}
+
+impl Node {
+    fn new(value: u32, cost: u32, children: Vec<usize>) -> Self {
+        Self {
+            value,
+            cost,
+            children,
+        }
+    }
+}
+
+pub fn run(input: &str, output: Option<&str>) -> Result<String> {
+    let mut lines = input.lines();
+    let first_line = split_whitespace_and_parse::<u32>(lines.next().unwrap())?;
+
+    let n = first_line[0];
+    let budget = first_line[1];
+
+    let mut tree = vec![Node::new(0, 0, vec![]); n as usize];
+
+    for _ in 0..n {
+        let parts = split_whitespace_and_parse::<u32>(lines.next().unwrap())?;
+        let id = parts[0] as usize - 1;
+        let value = parts[1];
+        let cost = parts[2];
+
+        tree[id] = Node::new(value, cost, vec![]);
+    }
+
+    for _ in 0..(n - 1) {
+        let parts = split_whitespace_and_parse::<usize>(lines.next().unwrap())?;
+        let parent_id = parts[0] - 1;
+        let child_id = parts[1] - 1;
+
+        tree[parent_id].children.push(child_id);
+    }
+
+    let result = max_value_within_budget(&tree, 0, budget);
+
+    if let Some(output) = output {
+        let expected = output.trim().parse::<u32>()?;
+
+        anyhow::ensure!(
+            result == expected,
+            "Output does not match expected: got '{result}' expected '{expected}'",
+        );
+
+        return Ok(format!("Result matches the specified output: {}", result));
+    }
+
+    Ok(format!(
+        "Calculated the maximum value: {result}, but {} output!",
+        "missing".red().bold()
+    ))
+}
+
+fn max_value_within_budget(tree: &[Node], root_id: usize, budget: u32) -> u32 {
+    let mut memo = vec![vec![None; budget as usize + 1]; tree.len()];
+    dfs(tree, root_id, budget, &mut memo)
+}
+
+fn dfs(tree: &[Node], node_id: usize, remaining_budget: u32, memo: &mut [Vec<Option<u32>>]) -> u32 {
+    let node = &tree[node_id];
+
+    // If we've already computed this state, return memoized result
+    if let Some(result) = memo[node_id][remaining_budget as usize] {
+        return result;
+    }
+
+    // If the node cost exceeds remaining budget, we can't keep this node or its subtree
+    if node.cost > remaining_budget {
+        memo[node_id][remaining_budget as usize] = Some(0);
+        return 0;
+    }
+
+    let budget_after_node = remaining_budget - node.cost;
+    let mut value_with_node = node.value;
+
+    // For each child, compute the maximum value we can get
+    for &child_id in &node.children {
+        let child_value = dfs(tree, child_id, budget_after_node, memo);
+        value_with_node += child_value;
+    }
+
+    memo[node_id][remaining_budget as usize] = Some(value_with_node);
+    value_with_node
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn sample_tree() -> (Vec<Node>, u32) {
+        let tree = vec![
+            Node::new(5, 3, vec![1, 2]), // Node 1 (0-indexed)
+            Node::new(2, 1, vec![]),     // Node 2 (0-indexed)
+            Node::new(8, 4, vec![]),     // Node 3 (0-indexed)
+        ];
+        (tree, 5)
+    }
+
+    #[test]
+    fn test_sample() {
+        let (tree, budget) = sample_tree();
+        assert_eq!(max_value_within_budget(&tree, 0, budget), 7);
+    }
+
+    #[test]
+    fn test_root_too_expensive() {
+        let tree = vec![Node::new(5, 10, vec![1]), Node::new(2, 1, vec![])];
+        assert_eq!(max_value_within_budget(&tree, 0, 5), 0);
+    }
+
+    #[test]
+    fn test_all_within_budget() {
+        let tree = vec![Node::new(1, 1, vec![1]), Node::new(2, 1, vec![])];
+        assert_eq!(max_value_within_budget(&tree, 0, 2), 3);
+    }
+
+    #[test]
+    fn test_large_budget() {
+        let tree = vec![
+            Node::new(10, 2, vec![1, 2]),
+            Node::new(20, 3, vec![]),
+            Node::new(30, 4, vec![]),
+        ];
+        assert_eq!(max_value_within_budget(&tree, 0, 9), 60);
+    }
+}

src/challenge/mod.rs

@@ -4,6 +4,7 @@ mod challenge_03;
 mod challenge_04;
 mod challenge_05;
 mod challenge_06;
+mod challenge_07;
 
 pub fn name(challenge: &str) -> Option<&'static str> {
     match challenge {
@@ -13,6 +14,7 @@ pub fn name(challenge: &str) -> Option<&'static str> {
         "04" => Some("Bananas! Bananas! Bananas!"),
         "05" => Some("Anagramic Subsequence"),
         "06" => Some("Research Lab Recorder"),
+        "07" => Some("Tree Pruning for Budget"),
         _ => None,
     }
 }
@@ -25,6 +27,7 @@ pub fn run(challenge: &str, input: &str, output: Option<&str>) -> anyhow::Result
         "04" => challenge_04::run(input, output),
         "05" => challenge_05::run(input, output),
         "06" => challenge_06::run(input, output),
+        "07" => challenge_07::run(input, output),
         _ => anyhow::bail!("Error: Unknown challenge '{challenge}'"),
     }
 }