2386.找出数组的第 K 大和

【LetMeFly】2386.找出数组的第 K 大和：逆向思维（小根堆）

力扣题目链接：https://leetcode.cn/problems/find-the-k-sum-of-an-array/

给你一个整数数组 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 。

 

提示：

• n == nums.length
• 1 <= n <= 105
• -109 <= nums[i] <= 109
• 1 <= k <= min(2000, 2n)

方法一：逆向思维（小根堆）

• 最大和怎么找？数组中正数之和即是。
• 第二大和怎么找？最大和 的基础上减去一个很小的正数或加上一个很小的负数。
• 第三大和怎么找？最大和/第二大和 的基础上减去一个很小的正数或加上一个很小的负数。
• …

不难发现除了初始的最大和，负数和其绝对值的正数结果是等价的。因此我们可以：

遍历一遍数组，若为正整数则累加，为负则变为正（取绝对值）。这样最大和就找到了。

又因为求的是“序列”和而不是“子数组”和，也就是说元素可以不连续。因此我们可以：

对处理后的全是非负数的数组从小到大排个序。

接下来的操作就是不断地从“最大和”的基础上，减去非负数组中的一些数了。每次减去地尽可能小，一共减去$k$次。也就是说：

我们只需要找到这个递增非负数组的前$k$个最小和即可。

假设这个处理后的数组是[1, 2, 3]，那么怎么生成它的$2^3$个子序列呢？

• 初始序列为空：[]
• 在空序列的基础上添加1得到[1]
• 在[1]的基础上添加2得到[1, 2]；或将最后的1替换成2得到[2]
• 在[2]基础上添加3得到[2, 3]；或直接替换最后的2成[3]
• 在[1, 2]的基础上添加3得到[1, 2, 3]；或直接替换最后的2成[1, 3]

这不是很适合最小堆吗？

堆中存放(当前和, 添加到了第几个元素），每次选当前和最小的出堆，即为下一个“最小序列和”。

• 时间复杂度$O(n\log n + k\log k)$，其中$n=len(nums)$
• 空间复杂度$O(k)$

AC代码

C++

/*
1 2 3

      1
    12  2
 123   23  3
*/

typedef long long ll;
typedef pair<ll, int> pli;
class Solution {
public:
    ll kSum(vector<int>& nums, int k) {
        ll sum = 0;
        for (int& t : nums) {
            if (t >= 0) {
                sum += t;
            }
            else {
                t = -t;
            }
        }
        sort(nums.begin(), nums.end());
        priority_queue<pli, vector<pli>, greater<pli>> pq;
        pq.push({nums[0], 0});
        ll toDesc = 0;
        for (int i = 1; i < k; i++) {
            auto [nowSum, th] = pq.top();
            toDesc = nowSum;
            pq.pop();
            if (th == nums.size() - 1) {
                continue;
            }
            pq.push({nowSum + nums[th + 1], th + 1});
            pq.push({nowSum - nums[th] + nums[th + 1], th + 1});
        }
        return sum - toDesc;
    }
};

Python

# from typing import List
# import heapq

class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        sum_ = 0
        for i in range(len(nums)):
            if nums[i] >= 0:
                sum_ += nums[i]
            else:
                nums[i] = -nums[i]
        nums.sort()
        pq = [(nums[0], 0)]
        toDesc = 0
        for _ in range(1, k):
            nowSum, th = heapq.heappop(pq)
            toDesc = nowSum
            if th == len(nums) - 1:
                continue
            heapq.heappush(pq, (nowSum + nums[th + 1], th + 1))
            heapq.heappush(pq, (nowSum - nums[th] + nums[th + 1], th + 1))
        return sum_ - toDesc

同步发文于CSDN和我的个人博客，原创不易，转载经作者同意后请附上原文链接哦~
Tisfy：https://letmefly.blog.csdn.net/article/details/136588803