How to Compute Average of Levels in Binary Tree?


Given a binary tree, compute the average values for the nodes on each level and return them in the form of an array. For example:

Input:

    3
   / \
  9  20
    /  \
   15   7

Output: [3, 14.5, 11]
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].

Level by Level Tree Traversal by BFS (Breadth First Search)

The Breadth First Search (BFS) algorithm allows us to travel the binary tree or other trees level by level. So with this in mind, we can have a unordered_map (hash table) to store the number of nodes corresponding to the depth of the tree.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == NULL) {
            return {};
        }
        queue<pair<TreeNode*, int>> Q;
        Q.push(make_pair(root, 0));
        unordered_map<int, int> counter;
        vector<double> result;
        while (!Q.empty()) {
            auto p = Q.front();
            Q.pop();
            auto depth = p.second;
            if (counter.find(depth) == counter.end()) {
                counter[depth] = 1; // first element in this level
            } else {
                counter[depth] ++;
            }
            if (result.size() <= depth) {
                result.push_back(p.first->val);
            } else {
                // update the average number
                result[depth] = (result[depth] * (counter[depth] - 1) + p.first->val) / counter[depth];
            }
            if (p.first->left) {
                Q.push(make_pair(p.first->left, depth + 1));
            }
            if (p.first->right) {
                Q.push(make_pair(p.first->right, depth + 1));
            }            
        }
        return result;
    }
};
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == NULL) {
            return {};
        }
        queue<pair<TreeNode*, int>> Q;
        Q.push(make_pair(root, 0));
        unordered_map<int, int> counter;
        vector<double> result;
        while (!Q.empty()) {
            auto p = Q.front();
            Q.pop();
            auto depth = p.second;
            if (counter.find(depth) == counter.end()) {
                counter[depth] = 1; // first element in this level
            } else {
                counter[depth] ++;
            }
            if (result.size() <= depth) {
                result.push_back(p.first->val);
            } else {
                // update the average number
                result[depth] = (result[depth] * (counter[depth] - 1) + p.first->val) / counter[depth];
            }
            if (p.first->left) {
                Q.push(make_pair(p.first->left, depth + 1));
            }
            if (p.first->right) {
                Q.push(make_pair(p.first->right, depth + 1));
            }            
        }
        return result;
    }
};

Based on similar idea, however, we can use a map in C++ which is an ordered map (sorted) then iterate over the map to compute the average for each level. During the expanding the children nodes in BFS, we need to push the levels (always the current depth plus one). The map’s key is the depth and the value is the vector of the nodes (values) in that depth.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == NULL) {
            return {};
        }
        map<int, vector<int>> data;  // level and their associate values
        vector<double> r;
        queue<pair<TreeNode*, int>> q;
        q.push(make_pair(root, 1));
        while (q.size() > 0) {
            auto p = q.front();
            q.pop();
            if (data.find(p.second) == data.end()) { 
                data[p.second] = { p.first->val };
            } else {
                data[p.second].push_back(p.first->val);
            }
            if (p.first->left != NULL) {
                q.push(make_pair(p.first->left, p.second + 1));
            }
            if (p.first->right != NULL) {
                q.push(make_pair(p.first->right, p.second + 1));
            }
        }
        for (const auto &n: data) {
            vector<int> x = n.second;
            int64_t sum = 0;
            for (int i = 0; i < x.size(); ++ i) {
                sum += x[i];
            }
            r.push_back(sum * 1.0 / x.size());
        }
        return r;
    }
};
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == NULL) {
            return {};
        }
        map<int, vector<int>> data;  // level and their associate values
        vector<double> r;
        queue<pair<TreeNode*, int>> q;
        q.push(make_pair(root, 1));
        while (q.size() > 0) {
            auto p = q.front();
            q.pop();
            if (data.find(p.second) == data.end()) { 
                data[p.second] = { p.first->val };
            } else {
                data[p.second].push_back(p.first->val);
            }
            if (p.first->left != NULL) {
                q.push(make_pair(p.first->left, p.second + 1));
            }
            if (p.first->right != NULL) {
                q.push(make_pair(p.first->right, p.second + 1));
            }
        }
        for (const auto &n: data) {
            vector<int> x = n.second;
            int64_t sum = 0;
            for (int i = 0; i < x.size(); ++ i) {
                sum += x[i];
            }
            r.push_back(sum * 1.0 / x.size());
        }
        return r;
    }
};

However, the following BFS is more concise. We are always expanding the child nodes of the current nodes in the queue, which all belong to the same level. Thus, without the need for extra memory, we can compute the average.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == nullptr) return {};
        queue<TreeNode*> Q;
        Q.push(root);
        vector<double> r;
        while (!Q.empty()) {
            int sz = Q.size();
            double sum = 0;
            for (size_t i = 0; i < sz; ++ i) { // expanding next level's nodes
                auto p = Q.front();
                Q.pop();
                if (p->left) {
                    Q.push(p->left);
                }
                if (p->right) {
                    Q.push(p->right);
                }
                sum += p->val;
            }
            r.push_back(sum / sz);
        }
        return r;
    }
};
class Solution {
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root == nullptr) return {};
        queue<TreeNode*> Q;
        Q.push(root);
        vector<double> r;
        while (!Q.empty()) {
            int sz = Q.size();
            double sum = 0;
            for (size_t i = 0; i < sz; ++ i) { // expanding next level's nodes
                auto p = Q.front();
                Q.pop();
                if (p->left) {
                    Q.push(p->left);
                }
                if (p->right) {
                    Q.push(p->right);
                }
                sum += p->val;
            }
            r.push_back(sum / sz);
        }
        return r;
    }
};

Both methods are O(N) as they need to travel all the nodes. And the first method has memory complexity O(logN) as the logN is the average depth of the binary tree with N nodes and the second method is O(1) memory complexity. In both cases, we don’t consider the result vector as an extra memory spending.

Using the algorithms that are discussed in this post, we can easily solve this: How to Find Largest Value in Each Tree Row/Level using BFS or DFS Algorithm?

The Python’s implmentation of Breadth First Search Algorithm to compute the average levels of binary tree: Teaching Kids Programming – Breadth First Search Algorithm to Compute Average of Levels in Binary Tree

–EOF (The Ultimate Computing & Technology Blog) —

GD Star Rating
loading...
861 words
Last Post: How to Design/Implement a Queue using Two Stacks?
Next Post: Installing a VPN on Your Router Is Easier Than You Might Have Thought

The Permanent URL is: How to Compute Average of Levels in Binary Tree?

Leave a Reply