LeetCode: 637. Average of Levels in Binary Tree

The Problem

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

Input:  [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].

Solution

We opted for a recursive approach to tackle this problem, and upon evaluating our solution on the LeetCode platform, we achieved the following outcome:

Here's the code that led us to this result.

vector<double> averageOfLevels(TreeNode* root) {
    vector<double> averages;
    vector<int> nodes;

    averageOfLevels (root, averages, nodes, 0);

    for(int i = 0; i < nodes.size(); ++i)
		averages[i] /= nodes[i];

    return averages;
}

void averageOfLevels ( TreeNode* root, vector<double> &averages, vector<int> &nodes, int level ) {
    if(root == nullptr) return;

    int nextLevel = level + 1; 

    if(averages.size() < nextLevel) {
    	averages.resize(nextLevel);
    	nodes.resize(nextLevel);
    }

    averages[level] += root->val;
    ++nodes[level];

    averageOfLevels (root->left, averages, nodes, nextLevel);
    averageOfLevels (root->right, averages, nodes, nextLevel);
}

Here's a step-by-step explanation:

The averageOfLevels function is a recursive helper function. It calculates the sum of node values and the count of nodes at each level of the binary tree. The function is called for each node in the tree.

Here's what it does:

1. It checks if the current root node is nullptr, indicating an empty subtree. If it is, the function simply returns.
2. It increments the nextLevel variable to represent the level one step deeper in the tree.
3. It checks whether the vectors averages and nodes need to be resized to accommodate the new level. If the current size is smaller than nextLevel, it resizes both vectors to the required size.
4. It adds the value of the current node (root->val) to the sum of node values at the current level (averages[level]), and increments the node count at the current level (nodes[level]).
5. It makes recursive calls to averageOfLevels for the left and right subtrees of the current node, passing in the updated vectors and the incremented nextLevel.