Leetcode.2583 二叉树中的第 K 大层和

题目链接

Leetcode.2583 二叉树中的第 K 大层和 Rating ： 1374

题目描述

给你一棵二叉树的根节点 root和一个正整数 k 。

树中的 层和 是指 同一层 上节点值的总和。

返回树中第 k 大的层和（不一定不同）。如果树少于 k 层，则返回 -1 。

注意，如果两个节点与根节点的距离相同，则认为它们在同一层。

示例 1：

输入：root = [5,8,9,2,1,3,7,4,6], k = 2
输出：13
解释：树中每一层的层和分别是：

• Level 1: 5
• Level 2: 8 + 9 = 17
• Level 3: 2 + 1 + 3 + 7 = 13
• Level 4: 4 + 6 = 10 第 2 大的层和等于 13 。

示例 2：

输入：root = [1,2,null,3], k = 1
输出：3
解释：最大的层和是 3 。

提示：

• 树中的节点数为 n
• $2<=n<=10^5$
• $1<=Node.val<=10^6$
• $1<=k<=n$

解法：bfs + 快速排序

先用 bfs ，把 层和 记录到数组 $ans$ 中。

再用 快速排序 ，求第 $k$ 大的层和。

时间复杂度： $O(n+nlogn)$

C++代码：

using LL = long long;

class Solution {
public:
    long long kthLargestLevelSum(TreeNode* root, int k) {
        vector<LL> ans;
        queue<TreeNode*> q;
        q.push(root);

        while(!q.empty()){
            int sz = q.size();
            LL s = 0;

            for(int i = 0;i < sz;i++){
                auto t = q.front();
                q.pop();

                s += t->val;
                if(t->left) q.push(t->left);
                if(t->right) q.push(t->right);
            }
            ans.push_back(s);
        }

        if(ans.size() < k) return -1;
        
        //把第 k 大的元素 放到位置 k-1 上
        nth_element(ans.begin(),ans.begin() + k - 1,ans.end(),greater{});


        return ans[k-1];
    }
};