题解 | #牛群平均重量# 层序遍历

知识点

层序遍历

思路

层序遍历得到每一层的val之和以及总个数。最后计算平均值。

时间复杂度

层序遍历，每个点只遍历一次，时间复杂度为

AC code （C++）

/**
 * struct TreeNode {
 *	int val;
 *	struct TreeNode *left;
 *	struct TreeNode *right;
 *	TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 * };
 */
class Solution {
public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 
     * @param root TreeNode类 
     * @return double浮点型vector
     */
    using pti = pair<TreeNode*, int>;
    vector<double> averageOfLevels(TreeNode* root) {
        vector<int> sum, cnt;
        queue<pti> q;
        if (root) q.emplace(root, 0);
        while (!q.empty()) {
            auto [p, idx] = q.front();
            q.pop();
            if (idx >= sum.size()) {
                sum.push_back(0);
                cnt.push_back(0);
            }
            sum[idx] += p->val;
            cnt[idx] += 1;
            if (p->left) q.emplace(p->left, idx + 1);
            if (p->right) q.emplace(p->right, idx + 1);
        }

        int m = sum.size();
        vector<double> res(m);
        for (int i = 0; i < m; i ++) {
            res[i] = 1.0 * sum[i] / cnt[i];
        }
        return res;
    }
};