description

Given the root of a binary tree, return the average value of the nodes on each level in the form of an array. Answers within 10-5 of the actual answer will be accepted.

 

Example 1:

Input: root = [3,9,20,null,null,15,7]
Output: [3.00000,14.50000,11.00000]
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].

Example 2:

Input: root = [3,9,20,15,7]
Output: [3.00000,14.50000,11.00000]

 

Constraints:

  • The number of nodes in the tree is in the range [1, 104].
  • -231 <= Node.val <= 231 - 1

submission

// Definition for a binary tree node.
// #[derive(Debug, PartialEq, Eq)]
// pub struct TreeNode {
//   pub val: i32,
//   pub left: Option<Rc<RefCell<TreeNode>>>,
//   pub right: Option<Rc<RefCell<TreeNode>>>,
// }
// 
// impl TreeNode {
//   #[inline]
//   pub fn new(val: i32) -> Self {
//     TreeNode {
//       val,
//       left: None,
//       right: None
//     }
//   }
// }
use std::rc::Rc;
use std::cell::RefCell;

impl Solution {
    // Many BFS problem can be solved with double queue, we alternate between
    // the two queues, iterate through one while adding new nodes to another
    // this is very elegant IMO
    // another advantage of this approach is we could use any container we like
    pub fn average_of_levels(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<f64> {
        let (mut dq, mut swp) = (Vec::new(), Vec::new());
        let mut ans = Vec::new();
        // setup, add the first node to queue
        if let Some(root) = root {
            dq.push(root);
        };
        // outer loop
        while !dq.is_empty() {
            // we setup the inner loop here, like sum
            let mut sum = 0.0f64;
            let cnt = dq.len() as f64;

            while let Some(node) = dq.pop() {
                let node = node.borrow();
                sum += node.val as f64;
                if let Some(ref n) = node.left {
                    swp.push(n.clone());
                }
                if let Some(ref n) = node.right {
                    swp.push(n.clone());
                }
            }
            // finish inner loop here, like calculate average
            ans.push(sum / cnt);
            // alternate between two queues
            std::mem::swap(&mut dq, &mut swp);
        }
        ans
    }
}