优先队列2386. 找出数组的第 K 大和

问题描述

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

问题求解

观察1: 如果求第一最大的子序列和是简单的，就是所有正数和；
观察2: 如果求正数序列第k大的子序列和是容易的，可以通过堆的方式进行高效求解；
观察3: 这里的负数其实可以转换为正数进行求解，加上负数等于减去对应的正数；

class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        n = len(nums)
        tot = 0
        for i, num in enumerate(nums):
            if num > 0: tot += num
            else: nums[i] = -num
        if k == 1: return tot
        nums.sort()
        h = [[nums[0], 0]]
        for _ in range(2, k):
            s, i = heappop(h)
            if i + 1 < n:
                heappush(h, [s + nums[i + 1], i + 1])
                heappush(h, [s - nums[i] + nums[i + 1], i + 1])
        return tot - heappop(h)[0]