↓Skip to main content
  1. Posts/

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

·2 mins·
Table of Contents

LeetCode-3086 拾起 K 个 1 需要的最少行动次数 #

Solution 1 #

首先, 获取 “1” 最容易的方式是直接选一个有 “1” 的坐标, 并且将左右侧的 “1” 移动过来. 对于后续的其它 “1”, 贪心地选择执行动作一与动作二是最优的. 如果动作一的次数足够多, 那么最少行动次数就是获取至多 3 个连续的 “1” + 后续 “1” 的数量乘以 2. 现在问题是, 如果动作一的次数不够, 怎么办? 此时我们需要不停使用动作二将远处的 “1” 移动过来. 注意到最开始获取 “1” 的方式可以合并到这种情况里, 因此不妨考虑先进行动作二取到 $k - maxChanges$ 个 “1” 的最小代价 (这是一个货仓选址问题, 把坐标选择在中位数处最优), 再加上后续使用动作一与动作二这一组合的代价. 代码如下:

class Solution:
    def minimumMoves(self, nums: List[int], k: int, maxChanges: int) -> int:
        n = len(nums)
        pos = []
        c = 0 # 统计数组中是否有连续 1, 2, 3 个 "1"
        for i, v in enumerate(nums):
            if v == 0:
                continue
            pos.append(i)
            c = max(c, 1)
            if i > 0 and nums[i - 1] == 1:
                if i > 1 and nums[i - 2] == 1:
                    c = 3
                else:
                    c = max(c, 2)
        c = min(c, k)
        # 如果 maxChanges 足够大
        if c + maxChanges >= k:
            return max(c - 1, 0) + (k - c) * 2

        # 如果 maxChanges 不够大
        idx = pos[len(pos) // 2]
        m = k - maxChanges # 从连续 m 个 "1" 中计算选在中位数的代价

        ans = inf

        presum = list(accumulate(pos, initial=0)) # presum[i] 为加到第 i - 1 个前缀和

        for r in range(m, len(pos) + 1):
            # [l, r)
            l = r - m

            mid = l + m // 2  # 中位数选址

            s1 = (mid - l) * pos[mid] - (presum[mid] - presum[l]) 
            s2= presum[r] - presum[mid] - (r - mid) * pos[mid]
            ans = min(ans, s1 + s2)
            
        return ans + 2 * maxChanges
↑