LC 6155. 找出数组的第 K 大和

https://leetcode.cn/problems/find-the-k-sum-of-an-array/

初看这题觉得这个数据量规模比较大，但是细看的话可以看到其中K最大值是2000, 所以必须从K入手。

然后就是可以转变一下这个问题：

假设我们有数组是 [1,2,3,4,5,-2,-3,-4,-5]
那么最大值肯定是 [1,2,3,4,5] = 15
假设就是选择从中减去 A=[1,2,3,4,5], 以及从中加上 B=[-2,-3,-4,-5]
而B其实可以变为从中减去 C=[2,3,4,5]
因为最后问题就变为，我们要从A+C中选择第K小的子序列之和，然后用最大值减去它。

从一个数组中选择第K小的子序列之和，大致思路有点类似Dijkstra贪心算法，进行K轮迭代每轮都要进行节点扩展，然后使用堆维护最小值。一个简单的优化是，对于第S轮，我们只需要扩展K-S个节点，这样常数项大致可以减半，否则会出现超时的情况。

class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        print(len(nums))
        a = [x for x in nums if x >= 0]
        base = sum(a)
        a += [-x for x in nums if x < 0]
        a.sort()

        import heapq
        hp = [(0, -1)]
        for step in range(k - 1):
            value, idx = heapq.heappop(hp)
            for i in range(k - step):
                j = idx + 1 + i
                if j < len(a):
                    heapq.heappush(hp, (value + a[j], j))
                else:
                    break
        ans = base - hp[0][0]
        return ans