Leetcode2583. 二叉树中的第 K 大层和

Every day a Leetcode

题目来源：2583. 二叉树中的第 K 大层和

解法1：层序遍历 + 排序

先使用层序遍历计算出树的每一层的节点值的和，保存在数组 levelSum 中。然后将数组进行排序，返回第 k 大的值。需要考虑数组长度小于 k 的边界情况。

代码：

/*
 * @lc app=leetcode.cn id=2583 lang=cpp
 *
 * [2583] 二叉树中的第 K 大层和
 */

// @lc code=start
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution
{
public:
    long long kthLargestLevelSum(TreeNode *root, int k)
    {
        if (root == nullptr)
            return -1;

        vector<long long> levelSum;
        queue<TreeNode *> q;
        q.push(root);
        while (!q.empty())
        {
            int size = q.size();
            long long sum = 0LL;
            for (int i = 0; i < size; i++)
            {
                TreeNode *node = q.front();
                q.pop();
                sum += node->val;
                if (node->left)
                    q.push(node->left);
                if (node->right)
                    q.push(node->right);
            }
            levelSum.push_back(sum);
        }

        if (levelSum.size() < k)
            return -1;
            
        sort(levelSum.begin(), levelSum.end());

        return levelSum[levelSum.size() - k];
    }
};
// @lc code=end

结果：

复杂度分析：

时间复杂度：O(nlogn)，其中 n 是二叉树的节点个数。

空间复杂度：O(n)，其中 n 是二叉树的节点个数。

解法2：层序遍历 + 快速选择

也可以使用快速选择的算法快速定位第 k 大的元素。

代码：

// 层序遍历 + 快速选择

class Solution
{
public:
    long long kthLargestLevelSum(TreeNode *root, int k)
    {
        if (root == nullptr)
            return -1;

        vector<long long> levelSum;
        queue<TreeNode *> q;
        q.push(root);
        while (!q.empty())
        {
            int size = q.size();
            long long sum = 0LL;
            for (int i = 0; i < size; i++)
            {
                TreeNode *node = q.front();
                q.pop();
                sum += node->val;
                if (node->left)
                    q.push(node->left);
                if (node->right)
                    q.push(node->right);
            }
            levelSum.push_back(sum);
        }

        int n = levelSum.size();
        if (k > n)
            return -1;
        ranges::nth_element(levelSum, levelSum.begin() + (n - k));
        return levelSum[n - k];
    }
};

结果：

复杂度分析：

时间复杂度：O(nlogn)，其中 n 是二叉树的节点个数。

空间复杂度：O(n)，其中 n 是二叉树的节点个数。