LeetCode 637.二叉树的层平均值

博主 913 2020-09-12

给定一个非空二叉树, 返回一个由每层节点平均值组成的数组。
eg:
输入:
3
/
9 20
/
15 7
输出:[3, 14.5, 11]
解释:
第 0 层的平均值是 3 , 第1层是 14.5 , 第2层是 11 。因此返回 [3, 14.5, 11] 。
提示:
节点值的范围在32位有符号整数范围内。

来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/average-of-levels-in-binary-tree
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。


思路很简单,就是维护两个数组,一个保存每层的节点数量,一个保存每层的节点总和。
可以使用dfs和bfs,我这里使用的dfs,因为bfs还需要转化一下层数的问题233
遍历了所有节点,时间复杂度为O(n),维护了两个大小为n的数组,空间复杂度为O(n)

代码:

/**
 * 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 {
private:
    vector<long long> sum;
    vector<int> count;
    vector<double> ans;
public:
    vector<double> averageOfLevels(TreeNode* root) {
        if (root != NULL) {
            dfs(root, 0);
        }
        // cout << sum.size() << "," << count.size()<< endl;
        for (int i = 0; i < sum.size(); ++i) {
            ans.push_back(sum[i]*1.0/count[i]);
        }
        return ans;
    }
    
    void dfs(TreeNode* root, int level){
//        cout << "当前val:" << root->val << endl;
//        cout << "sum size:"
        if (root != NULL) {
            if (sum.size() <= level) {
                sum.push_back(root->val);
                count.push_back(1);
            }else{
                sum[level] += root->val;
                ++count[level];
            }
            dfs(root->left, level+1);
            dfs(root->right, level+1);
        }
    }
};

结果:
image.png