637. Average of Levels in Binary Tree

Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.

Example 1:

Input:
    3
   / \
  9  20
    /  \
   15   7
Output: [3, 14.5, 11]
Explanation:
The average value of nodes on level 0 is 3,  on level 1 is 14.5, and on level 2 is 11. Hence return [3, 14.5, 11].

Note:

  1. The range of node's value is in the range of 32-bit signed integer.

Rust Solution

struct Solution;
use rustgym_util::*;

trait Preorder {
    fn preorder(&self, level: usize, levels: &mut Vec<Vec<i32>>);
}

impl Preorder for TreeLink {
    fn preorder(&self, level: usize, levels: &mut Vec<Vec<i32>>) {
        if let Some(node) = self {
            let node = node.borrow();
            let left = &node.left;
            let right = &node.right;
            if levels.len() == level {
                levels.push(vec![node.val]);
            } else {
                levels[level].push(node.val);
            }
            left.preorder(level + 1, levels);
            right.preorder(level + 1, levels);
        }
    }
}

impl Solution {
    fn average_of_levels(root: TreeLink) -> Vec<f64> {
        let mut levels: Vec<Vec<i32>> = vec![];
        root.preorder(0, &mut levels);
        levels
            .iter()
            .map(|v| v.iter().map(|&x| x as f64).sum::<f64>() as f64 / v.len() as f64)
            .collect()
    }
}

#[test]
fn test() {
    let root = tree!(3, tree!(9), tree!(20, tree!(15), tree!(7)));
    let res: Vec<f64> = vec![3f64, 14.5, 11f64];
    assert_eq!(Solution::average_of_levels(root), res);
}

Having problems with this solution? Click here to submit an issue on github.