3 637. Average of Levels in Binary Tree 🌟

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.

This question is in continuation with A general approach to level order traversal questions series.

Previous Questions

1. Binary tree level order traversal
2. Binary tree level order traversal - II
3. Binary tree zig-zag level order traversal

Recursive Solution

• Here we need to keep track of the sum of level and count of nodes so we can calculate the average.
• On every level we increment the level count and add the value of the node to the sum.
• If we encounter new level (level==levelSum.size()) we push new values(0.0) in the sum and count vectors.
• We do it recursively for left and right subtrees with increasing level count.
• Finally we can calculate average by dividing sum by count for each list.

Code

class Solution {
private:
    void averageOfLevelsDFS(vector<double>& levelSum, vector<double>& levelCount, TreeNode* root, int level)
    {
        if (root == NULL)
            return;
        if (level == levelSum.size()) {
            levelSum.push_back(0.0);
            levelCount.push_back(0.0);
        }

        levelSum[level] += root->val;
        levelCount[level] += 1;

        averageOfLevelsDFS(levelSum, levelCount, root->left, level + 1);
        averageOfLevelsDFS(levelSum, levelCount, root->right, level + 1);
    }

public:
    vector<double> averageOfLevels(TreeNode* root)
    {
        vector<double> levelSum, levelCount;

        averageOfLevelsDFS(levelSum, levelCount, root, 0);

        int n = levelSum.size();
        for (int i = 0; i < n; i++) {
            levelSum[i] = levelSum[i] / levelCount[i];
        }
        return levelSum;
    }
};

Iterative Solution

• This solution is so same as the first question of level order traversal.
• Here we just keep track of sum at current level instead of pushing node in level vector which we did in first question.
• At last we divide sum by size of the queue (i.e. count of nodes in current level).

class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root)
    {
        vector<double> res;
        if (root == NULL)
            return res;

        queue<TreeNode*> q;
        q.push(root);

        while (!q.empty()) {
            int sz = q.size();
            double sum = 0;
            for (int i = 0; i < sz; i++) {
                TreeNode* node = q.front(); q.pop();

                if (node->left != NULL) q.push(node->left);
                if (node->right != NULL) q.push(node->right);

                sum += node->val;
            }
            res.push_back(sum / sz);
        }
        return res;
    }
};

Iterative Solution

• Converted recursive solution to iterative.
• This solution involves 2 extra vectors which add up to the our additional use space.

code

class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root)
    {
        vector<double> levelSum, levelCount;

        queue<TreeNode*> q;
        q.push(root);

        while (!q.empty()) {
            int n = q.size();
            levelSum.push_back(0.0);
            levelCount.push_back(0.0);
            for (int i = 0; i < n; i++) {
                TreeNode* node = q.front(); q.pop();

                if (node->left) q.push(node->left);
                if (node->right) q.push(node->right);

                levelSum[levelSum.size() - 1] += node->val;
                levelCount[levelCount.size() - 1] += 1;
            }
        }
        int n = levelSum.size();
        for (int i = 0; i < n; i++) {
            levelSum[i] = levelSum[i] / levelCount[i];
        }
        return levelSum;
    }
};