最大子序列的分数

标签: 贪心 数组 排序 堆（优先队列）

难度: Medium

给你两个下标从 0 开始的整数数组 nums1 和 nums2 ，两者长度都是 n ，再给你一个正整数 k 。你必须从 nums1 中选一个长度为 k 的 子序列 对应的下标。

对于选择的下标 i0 ，i1 ，...， ik - 1 ，你的 分数 定义如下：

  • nums1 中下标对应元素求和，乘以 nums2 中下标对应元素的 最小值 。
  • 用公式表示： (nums1[i0] + nums1[i1] +...+ nums1[ik - 1]) * min(nums2[i0] , nums2[i1], ... ,nums2[ik - 1]) 。

请你返回 最大 可能的分数。

一个数组的 子序列 下标是集合 {0, 1, ..., n-1} 中删除若干元素得到的剩余集合，也可以不删除任何元素。

示例 1：

输入：nums1 = [1,3,3,2], nums2 = [2,1,3,4], k = 3
输出：12
解释：
四个可能的子序列分数为：
- 选择下标 0 ，1 和 2 ，得到分数 (1+3+3) * min(2,1,3) = 7 。
- 选择下标 0 ，1 和 3 ，得到分数 (1+3+2) * min(2,1,4) = 6 。
- 选择下标 0 ，2 和 3 ，得到分数 (1+3+2) * min(2,3,4) = 12 。
- 选择下标 1 ，2 和 3 ，得到分数 (3+3+2) * min(1,3,4) = 8 。
所以最大分数为 12 。

示例 2：

输入：nums1 = [4,2,3,1,1], nums2 = [7,5,10,9,6], k = 1
输出：30
解释：
选择下标 2 最优：nums1[2] * nums2[2] = 3 * 10 = 30 是最大可能分数。

提示：

  • n == nums1.length == nums2.length
  • 1 <= n <= 105
  • 0 <= nums1[i], nums2[j] <= 105
  • 1 <= k <= n

Submission

运行时间: 137 ms

内存: 34.4 MB

class Solution:
    def maxScore(self, nums1: List[int], nums2: List[int], K: int) -> int:
        # 按照 nums2 从小到大排序. 然后倒序遍历 zip. 选择 nums1 中最大的 k 个即可. 用 pq.
        # pq 怎么用？用 minPQ. 每次把最小的 pop 出来.
        n = len(nums1)
        arr = sorted(zip(nums1, nums2), key=lambda x: x[1])
        q = []
        # 倒着遍历：把最后 K 个 nums[n-k..n) 加进来.
        sm = 0
        for i in range(n-1, n-K-1, -1):
            sm += arr[i][0]
            heappush(q, arr[i][0])
        
        ans = sm * arr[n-K][1]
        
        # 接着倒着遍历.
        for i in range(n-K-1, -1, -1):
            sm += arr[i][0]
            heappush(q, arr[i][0])
            sm -= heappop(q)  # pop 过期的.
            ans = max(ans, sm * arr[i][1])
        
        return ans

Explain

题解的核心思路基于排序和贪心算法。首先，将数组 nums1 和 nums2 基于 nums2 的值进行排序。排序之后，对于每一个可能的最小值（即 nums2 中的元素），选择在这个最小值之后的 K 个 nums1 元素的最大可能和。这是通过使用一个最小堆（优先队列）来维护当前的最大 K 个元素来实现的。随着从后向前遍历排序后的数组，我们更新当前的和并检查是否可以通过替换堆中最小的元素来获得更大的值。这样可以保证每一步都使用了当前最小值之后的 K 个最大值来计算分数。

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

空间复杂度: O(n)

class Solution:
    def maxScore(self, nums1: List[int], nums2: List[int], K: int) -> int:
        # 按照 nums2 从小到大排序
        n = len(nums1)
        arr = sorted(zip(nums1, nums2), key=lambda x: x[1])
        q = []
        sm = 0
        # 从后向前遍历，保留 K 个最大的 nums1 元素
        for i in range(n-1, n-K-1, -1):
            sm += arr[i][0]
            heappush(q, arr[i][0])
        # 初始最大分数
        ans = sm * arr[n-K][1]
        # 继续向前遍历，更新最大分数
        for i in range(n-K-1, -1, -1):
            sm += arr[i][0]
            heappush(q, arr[i][0])
            sm -= heappop(q)  # 从堆中弹出最小值以保持堆大小为 K
            ans = max(ans, sm * arr[i][1])
        return ans

Explore

在题解中，数组nums1和nums2根据nums2的值进行排序的原因是：我们需要找到使得公式`(nums1[i0] + nums1[i1] +...+ nums1[ik - 1]) * min(nums2[i0], nums2[i1], ..., nums2[ik - 1])`的值最大化的子序列。因此，我们关注的是nums2数组中的最小值，因为它会与nums1子序列的和相乘。通过对nums2进行排序，我们可以更容易地控制和操作nums2的最小值。对于任意选择的最小值min(nums2)，我们希望nums1中相应位置的元素和尽可能大，因此这种排序方式可以有效地将问题转化为一个更易于解决的形式。

在题解中使用最小堆来确保我们始终能够维护K个最大元素的原因是，最小堆允许我们快速访问堆中的最小值。每当新的元素加入到堆中，如果堆的大小超过K，我们就将堆顶（即堆中的最小值）移除。这样，堆中始终保留的是当前遇到的最大的K个元素。通过这种方式，我们实际上是在利用最小堆的性质来确保堆中保存的是最大的K个值，而非随机选择的K个数。

题解中从后向前遍历数组的主要优势是方便地处理与最小值相关的逻辑。由于数组已经根据nums2进行排序，从后向前遍历可以保证我们每次处理的都是当前最小值之后的元素。这样可以简化逻辑，因为我们总是在减小的nums2值中选择最小值，这保证了每次计算的分数都是基于当前遍历到的位置确定的最小值。此外，这种方法还有助于我们在更新最大和时，更容易地通过比较新加入的元素和堆中的最小元素来优化总和。