Average of each level in a Binary Tree

Question

I am trying to find the average of each level in a binary tree. I am doing BFS. I am trying to do it using a null node. Whenever I find a dummy node, that means I am at the last node at that level. The problem I am facing is that I am not able to add average of the last level in a tree using this. Can Someone Help me?

Consider example [3,9,20,15,7] I am getting the output as [3.00000,14.50000]. Not getting the average of the last level that is 15 and 7 Here's my code

/**
* Definition for a binary tree node.
* public class TreeNode {
*     int val;
*     TreeNode left;
*     TreeNode right;
*     TreeNode(int x) { val = x; }
* }
*/

public class Solution {
 public List<Double> averageOfLevels(TreeNode root) {
    List<Double> list = new ArrayList<Double>();
    double sum = 0.0;
    Queue<TreeNode> q = new LinkedList<TreeNode>();
    TreeNode temp = new TreeNode(0);
    q.offer(root);
    q.offer(temp);
    int count = 0;
    while(!q.isEmpty()){
        root = q.poll();
        sum += root.val;

        if(root != temp)
        {
            count++;
            if(root.left != null){
                q.offer(root.left);
            }
            if(root.right != null){
                q.offer(root.right);
            }
        }
        else
        {
            if(!q.isEmpty()){
            list.add(sum / count);
            sum = 0;
            count = 0;
            q.add(temp);           
            }
      }

    }
    return list;
  }
}

Answer 1

I would use recursive deep scan of the tree. On each node I would push the value into a map with a pair .

I DID NOT test that code but it should be along the lines.

void scan(int level, TreeNode n, Map<Integer, List<Integer> m) {
  List l = m.get(level); if (l==null) {
    l = new ArrayList();
    m.put(level, l);
  }
  l.add(n.val);
  int nextLevel = level + 1;
  if (n.left != null) scan(nextLevel, n.left, m);
  if (n.right != null) scan(nextLevel, n.right, m);
}

Once the scan is done I can calculate the average for each level.

for (int lvl in m.keyset()) {
  List l = m.get(lvl);
  // MathUtils.avg() - it is obvious what it should be
  double avg = MathUtils.avg(l);
  // you code here
}

Answer 2

Take a look at this code, which executes whenever you find the marker for the end of the current level:

    if(!q.isEmpty()){
        list.add(sum / count);
        sum = 0;
        count = 0;
        q.add(temp);           
    }

This if statement seems to be designed to check whether you've finished the last row in the tree, which you could detect by noting that there are no more entries in the queue that would correspond to the next level. In that case, you're correct that you don't want to add the dummy node back into the queue (that would cause an infinite loop), but notice that you're also not computing the average in the row you just finished.

To fix this, you'll want to compute the average of the last row independently of reseeding the queue, like this:

   if(!q.isEmpty()){
        q.add(temp);           
   }

   list.add(sum / count);
   sum = 0;
   count = 0;

There's a new edge case to watch out for, and that's what happens if the tree is totally empty. I'll let you figure out how to proceed from here. Good luck!