Computer >> Computer tutorials >  >> Programming >> C++

Average of Levels in Binary Tree in C++


Suppose we have a non-empty binary tree; we have to find the average value of the nodes on each level in the return the average values as an array.

So, if the input is like

Average of Levels in Binary Tree in C++

then the output will be [3, 14.5, 11].

To solve this, we will follow these steps −

  • Define an array result

  • Define one queue q

  • insert root into q

  • while (not q is empty), do −

    • n := size of q

    • Define an array temp

    • while n is non-zero, do −

      • t := first element of q

      • insert value of t into temp

      • delete element from q

      • if left of t is not null, then −

        • insert left of t into q

      • if right of t is not null, then −

        • insert right of t into q

      • (decrease n by 1)

    • if size of temp is same as 1, then −

      • insert temp[0] at the end of result

    • otherwise when size of temp > 1, then −

      • sum := 0

      • for initialize i := 0, when i < size of temp, update (increase i by 1), do −

        • sum := sum + temp[i]

      • insert (sum / size of temp) at the end of result

  • return result

Example 

Let us see the following implementation to get a better understanding −

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
class TreeNode{
   public:
      int val;
      TreeNode *left, *right;
      TreeNode(int data){
         val = data;
         left = NULL;
         right = NULL;
      }
};
void insert(TreeNode **root, int val){
   queue<TreeNode*> q;
   q.push(*root);
   while(q.size()){
      TreeNode *temp = q.front();
      q.pop();
      if(!temp->left){
         if(val != NULL)
            temp->left = new TreeNode(val);
         else
            temp->left = new TreeNode(0);
         return;
      }
      else{
         q.push(temp->left);
      }
      if(!temp->right){
         if(val != NULL)
            temp->right = new TreeNode(val);
         else
            temp->right = new TreeNode(0);
         return;
      }
      else{
         q.push(temp->right);
      }
   }
}
TreeNode *make_tree(vector<int> v){
   TreeNode *root = new TreeNode(v[0]);
   for(int i = 1; i<v.size(); i++){
      insert(&root, v[i]);
   }
   return root;
}
class Solution{
public:
   vector<float> averageOfLevels(TreeNode *root){
      vector<float> result;
      queue<TreeNode*> q;
      q.push(root);
      while (!q.empty()) {
         int n = q.size();
         vector<float> temp;
         while (n) {
            TreeNode* t = q.front();
            temp.push_back(t->val);
            q.pop();
            if (t->left && t->left->val != 0)
               q.push(t->left);
            if (t->right && t->right->val != 0)
               q.push(t->right);
               n--;
         }
         if (temp.size() == 1)
            result.push_back(temp[0]);
         else if (temp.size() > 1) {
            double sum = 0;
            for (int i = 0; i < temp.size(); i++) {
               sum += temp[i];
            }
            result.push_back(sum / temp.size());
         }
      }
      return result;
   }
};
main(){
   Solution ob;
   vector<int> v = {3,9,20,NULL,NULL,15,7};
   TreeNode *root = make_tree(v);
   print_vector(ob.averageOfLevels(root));
}

Input

{3,9,20,NULL,NULL,15,7}

Output

[3, 14.5, 11, ]