题目链接
题目描述
给你一个整数数组 nums
和一个 正 整数 k
。你可以选择数组的任一 子序列 并且对其全部元素求和。
数组的 第 k 大和 定义为:可以获得的第 k
个 最大 子序列和(子序列和允许出现重复)
返回数组的 第 k 大和 。
子序列是一个可以由其他数组删除某些或不删除元素排生而来的数组,且派生过程不改变剩余元素的顺序。
注意:空子序列的和视作 0
。
示例 1:
输入:nums = [2,4,-2], k = 5 输出:2 解释:所有可能获得的子序列和列出如下,按递减顺序排列: - 6、4、4、2、2、0、0、-2 数组的第 5 大和是 2 。
示例 2:
输入:nums = [1,-2,3,4,-10,12], k = 16 输出:10 解释:数组的第 16 大和是 10 。
解题思路
计算 nums中所有非负数的和,记作 sum。
nums的任意一个子序列的元素和,都等价于从 sum中减去某些非负数 / 加上某些负数得到。
例如 nums=[1,2,3,−4],其非负数的和为 1+2+3=6,我们可以从6中减去2得到 nums的子序列 [1,3]的和 1+3=4,也可以把 6 和 −4 相加,得到 nums 的子序列 [1,2,3,−4] 的和 2。
注意到,「减去非负数」和「加上负数」都相当于减去 ∣nums[i]∣。sum减去的数越小,nums的子序列和就越大。
现在要解决的问题是:把每个 nums[i]取绝对值后,nums的第 k 小的子序列和是多少?
如何不重不漏地生成 nums的所有子序列?
以有序非负数组 nums=[1,2,3]为例,有2^3=8个子序列,生成的方法如下:
- 从空子序列 [] 开始。
- 在 [] 末尾添加 1 得到 [1]。
- 在 [1] 末尾添加 2 得到 [1,2]。也可以把末尾的 1 替换成 2 得到 [2]。
- 在 [2] 末尾添加 3 得到 [2,3]。也可以把末尾的 2 替换成 3 得到 [3]。
- 在 [1,2] 末尾添加 3 得到 [1,2,3]。也可以把末尾的 222 替换成 3 得到 [1,3]。
上述过程结合最小堆,就可以按照从小到大的顺序生成所有子序列了(堆中维护子序列的和以及下一个要添加/替换的元素下标)。取生成的第 k个子序列的和,作为 sum要减去的数。
复杂度分析
时间复杂度:O(nlogn+klogk),其中 n为 nums的长度。
空间复杂度:O(k)。
代码
#include <vector>
#include <iostream>
#include <queue>
#include <algorithm>
using namespace std;
class Solution {
public:
long long kSum(vector<int> &nums, int k) {
long sum = 0L;
for (int &x : nums) {
if (x >= 0) {
sum += x;
} else {
x = -x;
}
}
sort(nums.begin(), nums.end());
priority_queue<pair<long, int>, vector<pair<long, int>>, greater<>> pq;
pq.emplace(0, 0); // 空子序列
while (--k) {
auto [s, i] = pq.top();
pq.pop();
if (i < nums.size()) {
// 在子序列的末尾添加 nums[i]
pq.emplace(s + nums[i], i + 1); // 下一个添加/替换的元素下标为 i+1
if (i) { // 替换子序列的末尾元素为 nums[i]
pq.emplace(s + nums[i] - nums[i - 1], i + 1);
}
}
}
return sum - pq.top().first;
}
};
int main()
{
Solution solution;
vector<int> nums = {2, 4, -2};
int k = 5;
int res = solution.kSum(nums, k);
cout << res << endl;
nums = {1, -2, 3, 4, -10, 12};
k = 16;
res = solution.kSum(nums, k);
cout << res << endl;
}