拾起 K 个 1 需要的最少行动次数

难度: Hard

给你一个下标从 0 开始的二进制数组 nums，其长度为 n ；另给你一个 正整数 k 以及一个 非负整数 maxChanges 。

Alice 在玩一个游戏，游戏的目标是让 Alice 使用最少数量的行动次数从 nums 中拾起 k 个 1 。游戏开始时，Alice 可以选择数组 [0, n - 1] 范围内的任何索引 aliceIndex 站立。如果 nums[aliceIndex] == 1 ，Alice 会拾起一个 1 ，并且 nums[aliceIndex] 变成0（这不算作一次行动）。之后，Alice 可以执行 任意数量 的行动（包括零次），在每次行动中 Alice 必须恰好执行以下动作之一：

选择任意一个下标 j != aliceIndex 且满足 nums[j] == 0 ，然后将 nums[j] 设置为 1 。这个动作最多可以执行 maxChanges 次。
选择任意两个相邻的下标 x 和 y（|x - y| == 1）且满足 nums[x] == 1, nums[y] == 0 ，然后交换它们的值（将 nums[y] = 1 和 nums[x] = 0）。如果 y == aliceIndex，在这次行动后 Alice 拾起一个 1 ，并且 nums[y] 变成 0 。

返回 Alice 拾起恰好 k 个 1 所需的最少行动次数。

示例 1：

输入：nums = [1,1,0,0,0,1,1,0,0,1], k = 3, maxChanges = 1

输出：3

解释：如果游戏开始时 Alice 在 aliceIndex == 1 的位置上，按照以下步骤执行每个动作，他可以利用 3 次行动拾取 3 个 1 ：

游戏开始时 Alice 拾取了一个 1 ，nums[1] 变成了 0。此时 nums 变为 [1,1,1,0,0,1,1,0,0,1] 。
选择 j == 2 并执行第一种类型的动作。nums 变为 [1,0,1,0,0,1,1,0,0,1]
选择 x == 2 和 y == 1 ，并执行第二种类型的动作。nums 变为 [1,1,0,0,0,1,1,0,0,1] 。由于 y == aliceIndex，Alice 拾取了一个 1 ，nums 变为 [1,0,0,0,0,1,1,0,0,1] 。
选择 x == 0 和 y == 1 ，并执行第二种类型的动作。nums 变为 [0,1,0,0,0,1,1,0,0,1] 。由于 y == aliceIndex，Alice 拾取了一个 1 ，nums 变为 [0,0,0,0,0,1,1,0,0,1] 。

请注意，Alice 也可能执行其他的 3 次行动序列达成拾取 3 个 1 。

示例 2：

输入：nums = [0,0,0,0], k = 2, maxChanges = 3

输出：4

解释：如果游戏开始时 Alice 在 aliceIndex == 0 的位置上，按照以下步骤执行每个动作，他可以利用 4 次行动拾取 2 个 1 ：

选择 j == 1 并执行第一种类型的动作。nums 变为 [0,1,0,0] 。
选择 x == 1 和 y == 0 ，并执行第二种类型的动作。nums 变为 [1,0,0,0] 。由于 y == aliceIndex，Alice 拾起了一个 1 ，nums 变为 [0,0,0,0] 。
再次选择 j == 1 并执行第一种类型的动作。nums 变为 [0,1,0,0] 。
再次选择 x == 1 和 y == 0 ，并执行第二种类型的动作。nums 变为 [1,0,0,0] 。由于y == aliceIndex，Alice 拾起了一个 1 ，nums 变为 [0,0,0,0] 。

提示：

2 <= n <= 10⁵
0 <= nums[i] <= 1
1 <= k <= 10⁵
0 <= maxChanges <= 10⁵
maxChanges + sum(nums) >= k

Submission

运行时间: 112 ms

内存: 27.2 MB

class Solution:
    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int:
        c = 0  # 对于连续的3个位置，最多有几个连续1
        n = len(nums)
        for i in range(n):
            s = sum(nums[i: min(n, i + 3)])
            if s == 3:
                c = 3
                break
            if s == 2:
                if nums[i + 1] != 0:
                    c = max(c, 2)
                else:
                    c = max(c, 1)
            elif s == 1:
                c = max(c, 1)
        if c >= k:
            return k - 1
        if maxChanges + c >= k:
            return max(0, c - 1) + (k - c) * 2

        # 剩下的又要用到第二种操作
        # 优先用的是maxChange，剩下的就变成货仓选址问题，计算过程自然会对连续的3个位置进行不超过2次拾起全部1的操作
        # 因此不需要先进行连续3个位置的计算
        idx = [i for i, x in enumerate(nums) if x]
        s = list(accumulate(idx, initial=0))
        m = len(idx)
        ans = inf
        # 在idx中选连续 k - maxChange 个 1，将其都移到中位数处
        for i in range(m - (k - maxChanges) + 1):
            # 子数组 idx[i: i + k - maxChanges]
            if (k - maxChanges) & 1:  # 长度为奇数
                mid = i + (k - maxChanges) // 2
                s1 = s[mid] - s[i]
                s2 = s[i + k - maxChanges] - s[mid + 1]
            else:
                mid = i + (k - maxChanges) // 2
                s1 = s[mid] - s[i]
                s2 = s[i + k - maxChanges] - s[mid]
            ans = min(ans, s2 - s1)
        return ans + maxChanges * 2

Explain

这道题的解题思路可以分为几个部分。首先，检查连续三个位置中最多有几个连续的1，并将其存储在变量c中。如果c大于等于k，直接返回k-1，因为Alice可以直接在连续的1中拾取。接下来，如果maxChanges加上c大于等于k，意味着Alice可以通过改变0为1和移动1来拾取k个1，每个1需要两次操作，因此返回c-1加上(k-c)*2。如果上述情况都不满足，说明Alice需要使用第二种操作（移动1）来拾取剩余的1，这时问题转化为一个货仓选址问题，即在nums的1的位置数组idx中选择连续的k-maxChanges个1，并将它们移动到中位数位置，使得总移动距离最小。最后返回这个最小移动距离加上maxChanges*2。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int:
        c = 0  # 对于连续的3个位置，最多有几个连续1
        n = len(nums)
        for i in range(n):
            s = sum(nums[i: min(n, i + 3)])
            if s == 3:
                c = 3
                break
            if s == 2:
                if nums[i + 1] != 0:
                    c = max(c, 2)
                else:
                    c = max(c, 1)
            elif s == 1:
                c = max(c, 1)
        if c >= k:
            return k - 1
        if maxChanges + c >= k:
            return max(0, c - 1) + (k - c) * 2
        idx = [i for i, x in enumerate(nums) if x]  # 存储nums中1的位置
        s = list(accumulate(idx, initial=0))  # 计算前缀和
        m = len(idx)
        ans = inf
        for i in range(m - (k - maxChanges) + 1):
            if (k - maxChanges) & 1:  # 长度为奇数
                mid = i + (k - maxChanges) // 2
                s1 = s[mid] - s[i]
                s2 = s[i + k - maxChanges] - s[mid + 1]
            else:
                mid = i + (k - maxChanges) // 2
                s1 = s[mid] - s[i]
                s2 = s[i + k - maxChanges] - s[mid]
            ans = min(ans, s2 - s1)
        return ans + maxChanges * 2  # 返回最小移动距离加上maxChanges*2

Explore

选择长度为3的窗口主要是基于优化算法的性能和简化问题的处理。在考虑拾取1的最优策略时，较小的窗口（如3个位置）可以有效地评估局部最优（即连续的1的局部最大数量），从而快速决定是否可以通过简单操作（如直接拾取或少量修改）达成目标。此外，3个位置的窗口也足够捕捉到问题中的关键信息，而不至于过度复杂化算法。

当连续的1的数量c大于等于k时，Alice可以在这些连续的1中直接拾取所需的k个1，不需要任何额外的移动。因此，理论上拾取动作只需要进行k次。但是，题目可能隐含了Alice在拾取过程中的位置调整，每次从一个1跳到下一个1视为一次行动，这样总共需要k-1次跳跃。因此，在连续的1足够多时，返回k-1表示Alice在拾取完这k个1后的最小行动次数。

使用前缀和数组能够快速计算任意子数组的和，这在处理连续元素的移动和位置调整问题时尤其有效。在该问题中，需要计算特定长度的子数组（即连续的1的位置）中的元素向中位数移动的总距离。通过前缀和，可以快速获取任意区间内1的位置和，从而计算出移动到中位数的成本。这样就可以遍历所有可能的子数组，找到移动成本最小的配置，显著减少了计算的复杂性和时间消耗。