找出数组的第 K 大和

标签: 数组 排序 堆(优先队列)

难度: Hard

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

Submission

运行时间: 144 ms

内存: 28.6 MB

class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        n = len(nums)
        acc = sum(num for num in nums if num > 0)
        if k == 1: return acc

        # queue = sorted(abs(num) for num in nums)
        queue = heapq.nsmallest(k, list(abs(num) for num in nums))

        heap = [(queue[0], 0)]

        for _ in range(k-2):
            res, index = heappop(heap)
            next = index + 1
            if next == n: continue
            heappush(heap, (res - queue[index] + queue[next], next))
            heappush(heap, (res + queue[next], next))
        
        return acc - heappop(heap)[0]

Explain

这道题目的核心思路是利用最小堆来维护所有可能的子序列和,从而找出第 k 大的和。首先,计算所有正数的和 `acc`,这代表了最大的子序列和(即所有正数相加)。然后,使用一个最小堆来维护和操作。堆中的元素是以绝对值排序的,这是因为我们需要逐步从最大的子序列和中减去一些数值来获得次大的和。通过维护堆来迭代地计算可能的子序列和,并逐步获得第 k 大的和。代码中的每一步操作都是基于这个最小堆来进行的,我们通过弹出和推入新的元素来控制和更新堆的状态。

时间复杂度: O(n log k + k log k)

空间复杂度: O(k)

class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        n = len(nums)
        acc = sum(num for num in nums if num > 0)  # 计算所有正数的和
        if k == 1: return acc  # 如果 k 为 1,直接返回最大子序列和

        # 获取前 k 个最小的绝对值
        queue = heapq.nsmallest(k, list(abs(num) for num in nums))

        heap = [(queue[0], 0)]  # 初始化堆

        for _ in range(k-2):
            res, index = heappop(heap)  # 弹出堆顶元素
            next_index = index + 1
            if next_index == n: continue
            heappush(heap, (res - queue[index] + queue[next_index], next_index))  # 更新堆
            heappush(heap, (res + queue[next_index], next_index))  # 添加新可能的子序列和
        
        return acc - heappop(heap)[0]  # 计算第 k 大的和

Explore

在处理子序列和的问题中,选择绝对值排序而非原始数值排序的主要原因是为了确保我们能从最大和中逐步减去的数值是最小的,即逐步探索较小的子序列和。如果我们按原始数值排序,那么负数的处理会变得复杂,因为我们需要考虑加上一个负数实际上是减小总和。而通过绝对值排序,我们可以更直观地按数值大小处理各个数,从而有效地通过最小堆来控制和更新子序列和的大小。

在堆中存储元素时包含索引是为了追踪当前处理到哪个元素,以及从哪个元素派生出当前的子序列和。这里的索引表示的是当前元素在按绝对值排序后的数组中的位置。通过记录索引,我们可以确保每次操作都是基于正确的元素进行,避免重复处理同一个元素,同时也能按顺序探索所有可能的子序列和。这是一种有效的方式来维护和更新堆中的状态,确保我们能正确地找到第 k 大的子序列和。

在代码中直接返回累加的正数和`acc`如果k等于1确实是一个问题,尤其是当数组中所有元素都是负数时。这种情况下,最大的子序列和应该是0(即不选择任何元素),而不是负数的累加和。因此,代码在处理全负数数组时应该有额外的检查,以确保当k为1且所有数都是负数时返回0。这种情况下的直接返回`acc`实际上是忽略了全负数的特殊情况。

如果数组中包含重复元素,使用最小堆来维护和操作子序列和的方法仍然是有效的。重复元素在处理时会被视为独立的实体,因此在计算子序列和时,每个重复的元素都可以被单独考虑。这确保了算法的正确性和完整性,不会因为元素的重复而影响最终的结果。维护最小堆的过程考虑了所有可能的子序列和,包括由重复元素形成的所有组合。