题目
- 题目链接:https://leetcode.cn/problems/find-the-k-sum-of-an-array/
给你一个整数数组 nums(有正有负)和一个正整数 k 。你可以选择数组的任一 子序列 并且对其全部元素求和。
数组的 第 k 大和 定义为:可以获得的第 k 个 最大 子序列和(子序列和允许出现重复)
返回数组的 第 k 大和 。
解法
这道题还是挺有意思的一道题,其核心问题在于:如何快速的找出第 k 大的子序列。
一个比较直观的想法是用 DFS 不断去遍历(选/不选),这样复杂度直接上天,显然没法快速解决,于是开始找规律。
小根堆求第 k 小的子序列
std 解法比较 trick ,我也很好奇是如何想出这种解法的。首先将求第 k 大的子序列转化为求第 k 小的子序列和,然后用所有元素的和减去第k小子序列的和即可得到 第 k 大的子序列和。用大写的 SUM 表示所有元素的和。
- 将 nums 数组从小到大排序
- 创建一个小根堆,维护 <sum,idx> ,表示当前 nums[0]~nums[idx] 选出来的子序列为 sum
- 每次弹出堆顶元素,然后做以下两个操作:
- 选 idx 位置的元素,压入 <sum+nums[idx+1],idx+1>
- 不选 idx 位置的元素,压入 <sum-nums[idx]+nums[idx+1],idx+1>
- 弹出 k-1 次后,堆顶元素即为第 k 小的子序列和,用 SUM 减去即可。
这样做的证法,一个比较好的解释在 :https://www.bilibili.com/video/BV1md4y1P75q?vd_source=3a233bcf5ea77820111ff3d44b807551 链接的第 21 分钟。
简单的说就是每次对于idx的位置,都可以有选和不选的两种方法,后续一定是将 nums[idx] 换成 nums[idx+1] ,这样保证了新变成的子序列一定是所有比当前子序列大的子序列中最小的一个。
负数转化
因为题目中的 nums 可能包含负数,所以我们也要对负数做相应处理。
同样的,我们记所有正数的和为 SUM,然后将所有负数取绝对值排序,接着用上述思路求第 k 小,此时
- 如果当前数是正数,那么最后用 SUM 减的时候,就相当于减去了一个正数(类似于考试100分,做错了扣分但是不会帮你扣成负数)
- 如果当前数是负数,因为 SUM 本身就没有包括这个数,就相当于直接加上了这个负数(考试100分,能帮你把分数扣成负数)
最终的题解如下
class Solution {
public:
struct Node {
long long sum;
int idx;
bool operator<(const Node& other) const {
return sum > other.sum;
}
};
long long kSum(vector<int>& a, int k) {
int n = a.size();
long long s = 0;
long long neg = 0;
long long ans = 0;
for(int i=0;i<n;i++) {
if(a[i]>0) s+=a[i];
else {
a[i] = -a[i];
}
}
sort(a.begin(),a.end());
// for(auto i:a) cout << i << ' ';
// cout << endl;
priority_queue<Node> q;
q.push({a[0],0});
for(int i=0;i<k-1;i++) {
Node now = q.top();
q.pop();
// cout << now.sum << ' ' << now.idx << endl;
ans = now.sum;
if(now.idx==n-1)
continue;
q.push({now.sum+a[now.idx+1],now.idx+1});
q.push({now.sum-a[now.idx]+a[now.idx+1],now.idx+1});
}
// cout << s << ' ' << q.top().sum << endl;
return s - ans;
}
};