LeetCode 637.二叉树的层平均值

给定一个非空二叉树, 返回一个由每层节点平均值组成的数组。
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);
        }
    }
};

结果：