2386. 找出数组的第 K 大和

2022/08/27 leetcode 共 1452 字,约 5 分钟

title

problem link

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

solution

参考

要记住

  • 数组有序后,可以用堆动态生成子序列的和,按递增顺序,每取一个后便加下一个,同时当前也可减掉。
  • 回溯只要控制深度复杂度也可观,可用于二分等

优先队列做法


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

二分做法

class BinarySearch:
    # If you wanna binary search big integer, plz set data range, which can be used as same as big integer low and high
    # find the first index that value >= val
    def bisect_left(data, val, key=None):
        l, r = 0, len(data) - 1
        if key is None:
            key = lambda x: data[mid]
        while l <= r:
            mid = (l + r) // 2
            if key(mid) >= val:
                r = mid - 1
            else:
                l = mid + 1
        return l
class Solution:
    def kSum(self, nums: List[int], k: int) -> int:
        s = 0
        for i, x in enumerate(nums):
            if x >= 0:
                s += x
            else:
                nums[i] = -x
        total = sum(nums)
        nums.sort()

        # 不超过limit的子序列数量
        def count(limit):
            cnt = 0

            # 算i之后的子序列数量,不包括空序列
            def f(cur, i):
                nonlocal cnt
                if cnt >= k - 1 or i >= len(nums) or cur + nums[i] > limit:
                    return
                cnt += 1
                f(cur + nums[i], i + 1)
                f(cur, i + 1)

            f(0, 0)
            return cnt

        j = BinarySearch.bisect_left(range(total), k - 1, key=count)

        return s - j

Search

    Table of Contents