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 是二叉树的节点个数。