幸福值最大化的选择方案

标签: 贪心 数组 排序

难度: Medium

给你一个长度为 n 的数组 happiness ，以及一个 正整数 k 。

n 个孩子站成一队，其中第 i 个孩子的 幸福值 是 happiness[i] 。你计划组织 k 轮筛选从这 n 个孩子中选出 k 个孩子。

在每一轮选择一个孩子时，所有 尚未 被选中的孩子的 幸福值 将减少 1 。注意，幸福值 不能 变成负数，且只有在它是正数的情况下才会减少。

选择 k 个孩子，并使你选中的孩子幸福值之和最大，返回你能够得到的 最大值 。

示例 1：

输入：happiness = [1,2,3], k = 2
输出：4
解释：按以下方式选择 2 个孩子：
- 选择幸福值为 3 的孩子。剩余孩子的幸福值变为 [0,1] 。
- 选择幸福值为 1 的孩子。剩余孩子的幸福值变为 [0] 。注意幸福值不能小于 0 。
所选孩子的幸福值之和为 3 + 1 = 4 。

示例 2：

输入：happiness = [1,1,1,1], k = 2
输出：1
解释：按以下方式选择 2 个孩子：
- 选择幸福值为 1 的任意一个孩子。剩余孩子的幸福值变为 [0,0,0] 。
- 选择幸福值为 0 的孩子。剩余孩子的幸福值变为 [0,0] 。
所选孩子的幸福值之和为 1 + 0 = 1 。

示例 3：

输入：happiness = [2,3,4,5], k = 1
输出：5
解释：按以下方式选择 1 个孩子：
- 选择幸福值为 5 的孩子。剩余孩子的幸福值变为 [1,2,3] 。
所选孩子的幸福值之和为 5 。

提示：

  • 1 <= n == happiness.length <= 2 * 105
  • 1 <= happiness[i] <= 108
  • 1 <= k <= n

Submission

运行时间: 144 ms

内存: 39.7 MB

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        happiness.sort(reverse = True)
        if happiness[k-1] >= k:
            return sum(happiness[:k]) - (k* (k-1))//2
        t = 0
        for i in range(k):
            if happiness[i] >= t:
                t += 1
            else:
                break
        return sum(happiness[:t]) - (t* (t-1))//2

Explain

首先，题解中对happiness数组进行了降序排序，这样可以直接从最大的幸福值开始选择。接下来，会检查第k个元素是否大于等于k，如果是，直接计算前k个元素的和减去一个特定的序列和（0到k-1的和），这是因为每次选择都会使得未选择的元素幸福值减1。如果第k个元素小于k，那么需要寻找一个最大的t（t<=k），使得每个happiness[i]都至少为t，此时计算前t个元素的和减去序列和。这种方法确保了在每轮减少幸福值的情况下，能够获得最大的幸福值总和。

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

空间复杂度: O(1)

class Solution:
    def maximumHappinessSum(self, happiness: List[int], k: int) -> int:
        happiness.sort(reverse=True) # 将幸福值数组降序排序
        if happiness[k-1] >= k:
            # 如果第k个元素的幸福值>=k，可以直接计算前k个元素的和减去(0到k-1)的和
            return sum(happiness[:k]) - (k * (k-1)) // 2
        t = 0
        for i in range(k):
            if happiness[i] >= t:
                t += 1 # t是可以持续选取的孩子的数量，直到happiness[i]<t
            else:
                break
        return sum(happiness[:t]) - (t * (t-1)) // 2 # 返回前t个元素的和减去(0到t-1)的和

Explore

在题解中，t代表可以连续选择的孩子的数量。初始化t为0是因为开始时没有孩子被选择。通过逐个检查每个孩子的幸福值是否大于等于t，可以确保每次选择后，剩余的每个孩子的幸福值至少减少1，而且保证选择的孩子数量不超过他们原始的幸福值。这种方法保证了在每一轮选择中，我们都取得了当前可能的最大幸福值，因为每增加一个孩子，我们都确保了这个孩子的幸福值在减少后仍然是正数，从而最大化了总幸福值。

如果happiness[k-1]（即第k个最大的幸福值）大于等于k，这意味着前k个孩子的幸福值都至少为k。因此，即使每选择一次，所有孩子的幸福值减少1，这些孩子的幸福值在k轮选择后仍然为非负。这样，选择前k个孩子提供的总幸福值最大，且每个孩子至少可以被选择一次，不会因为幸福值的减少而影响到总选择次数。因此，在这种情况下我们可以安全地直接计算这k个孩子的幸福值和减去序列和，而不必担心后续幸福值的减少问题。

当happiness[k-1]（即第k个最大的幸福值）小于k时，提示我们不能简单地选择前k个孩子，因为至少有一个孩子的幸福值少于k，意味着在连续选择k次后，这个孩子的幸福值会变为负数，违反了选择规则。此时，我们需要重新确定最大的t值，即最大的连续可选择孩子数，使得每个孩子的幸福值在选择过程中始终不少于选择次数。通过从头开始检查每个孩子的幸福值是否至少为当前的选择次数（t），并逐步增加t，直到找到使得happiness[t-1] < t的点为止。这样，我们可以确保在选择t个孩子的过程中，每个孩子的幸福值始终满足选择条件，从而获得最大的总幸福值。