algorithms_edu/algo/graph/tree/
sum.rs

1//! # Tree Sum Example
2//!
3//! - Time Complexity: O(n)
4//!
5//! # Resources
6//!
7//! - [W. Fiset's video](https://www.youtube.com/watch?v=0qgaIMqOEVs&list=PLDV1Zeh2NRsDGO4--qE8yH72HFL1Km93P&index=8)
8
9// making the tree generic over all summable types, e.g. `f64`, `i32`, or your own
10// `Complex` type (as long as it implements these traits)
11pub trait Summable: std::ops::AddAssign<Self> + Copy + num_traits::Zero {}
12impl<T: std::ops::AddAssign<Self> + Copy + num_traits::Zero> Summable for T {}
13
14pub struct TreeNode<T: Summable> {
15    val: T,
16    children: Vec<Box<TreeNode<T>>>,
17}
18
19impl<T: Summable> TreeNode<T> {
20    pub fn new(val: T) -> Self {
21        Self {
22            val,
23            children: Vec::new(),
24        }
25    }
26    pub fn add_child(&mut self, child: TreeNode<T>) {
27        self.children.push(Box::new(child));
28    }
29
30    pub fn sum(&self) -> T {
31        self.children
32            .iter()
33            .fold(T::zero(), |sum, child| sum + child.sum())
34            + self.val
35    }
36    pub fn leaf_sum(&self) -> T {
37        // a leaf has no children
38        if self.children.is_empty() {
39            self.val
40        } else {
41            self.children
42                .iter()
43                .fold(T::zero(), |sum, child| sum + child.leaf_sum())
44        }
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn test_tree_sum() {
54        let mut root = TreeNode::new(5);
55        let mut node4 = TreeNode::new(4);
56        let mut node3 = TreeNode::new(3);
57        let mut node1 = TreeNode::new(1);
58        let mut node7 = TreeNode::new(7);
59
60        let nodem6 = TreeNode::new(-6);
61        let node0 = TreeNode::new(0);
62        let nodem4 = TreeNode::new(-4);
63        let node2 = TreeNode::new(2);
64        let node9 = TreeNode::new(9);
65        let node8 = TreeNode::new(8);
66
67        node1.add_child(node2);
68        node1.add_child(node9);
69        node4.add_child(node1);
70        node4.add_child(nodem6);
71        root.add_child(node4);
72        node3.add_child(node0);
73        node7.add_child(node8);
74        node3.add_child(node7);
75        node3.add_child(nodem4);
76        root.add_child(node3);
77
78        let sum = root.sum();
79        println!("Tree sum: {}", sum);
80        assert_eq!(sum, 29);
81
82        let leaf_sum = root.leaf_sum();
83        println!("Leaf sum: {}", leaf_sum);
84        assert_eq!(leaf_sum, 9);
85    }
86}
algorithms_edu/algo/graph/tree/sum.rs

algorithms_edu/algo/graph/tree/
sum.rs
