rustgym/leetcode/
_637_average_of_levels_in_binary_tree.rs

1struct Solution;
2use rustgym_util::*;
3
4trait Preorder {
5    fn preorder(&self, level: usize, levels: &mut Vec<Vec<i32>>);
6}
7
8impl Preorder for TreeLink {
9    fn preorder(&self, level: usize, levels: &mut Vec<Vec<i32>>) {
10        if let Some(node) = self {
11            let node = node.borrow();
12            let left = &node.left;
13            let right = &node.right;
14            if levels.len() == level {
15                levels.push(vec![node.val]);
16            } else {
17                levels[level].push(node.val);
18            }
19            left.preorder(level + 1, levels);
20            right.preorder(level + 1, levels);
21        }
22    }
23}
24
25impl Solution {
26    fn average_of_levels(root: TreeLink) -> Vec<f64> {
27        let mut levels: Vec<Vec<i32>> = vec![];
28        root.preorder(0, &mut levels);
29        levels
30            .iter()
31            .map(|v| v.iter().map(|&x| x as f64).sum::<f64>() as f64 / v.len() as f64)
32            .collect()
33    }
34}
35
36#[test]
37fn test() {
38    let root = tree!(3, tree!(9), tree!(20, tree!(15), tree!(7)));
39    let res: Vec<f64> = vec![3f64, 14.5, 11f64];
40    assert_eq!(Solution::average_of_levels(root), res);
41}