[LeetCode]637. Average of Levels in Binary Tree

[LeetCode]637. Average of Levels in Binary Tree

题目描述

思路

实际上是树的层次遍历
递归写法参数中加入层数即可

代码

#include <iostream>
#include <vector>

using namespace std;

struct TreeNode {
    int val;
    TreeNode* left;
    TreeNode* right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution {
public:
    void averageOfLevels(TreeNode* root, int level, vector<double>& sum, vector<double>& count) {
        if (root == NULL) return;
        if (level < sum.size()) {
            sum[level] += root->val;
            count[level] += 1;
        }
        else {
            sum.push_back(root->val);
            count.push_back(1);
        }
        averageOfLevels(root->left, level + 1, sum, count);
        averageOfLevels(root->right, level + 1, sum, count);
    }


    vector<double> averageOfLevels(TreeNode* root) {
        vector<double> sum, count;
        averageOfLevels(root, 0, sum, count);
        vector<double> res(sum.size(), 0);
        for (int i = 0; i < res.size(); ++i)
            res[i] = sum[i] / count[i];
        return res;
    }
};


int main() {
    vector<double> res;
    TreeNode* t1 = new TreeNode(3);
    TreeNode* t2 = new TreeNode(9);
    TreeNode* t3 = new TreeNode(20);
    TreeNode* t4 = new TreeNode(15);
    TreeNode* t5 = new TreeNode(7);

    t1->left = t2;
    t1->right = t3;

    t2->left = t4;
    t3->right = t5;

    Solution s;
    res = s.averageOfLevels(t1);
    for (double num : res)
        cout << num << ' ';
    cout << endl;

    system("pause");
}