1703. 得到连续 K 个 1 的最少相邻交换次数 #

问题描述#

给你一个整数数组 nums 和一个整数 k 。 nums 仅包含 0 和 1 。每一次移动，你可以选择相邻两个数字并将它们交换。

请你返回使 nums 中包含 k 个连续 1 的最少交换次数。

示例 1：
输入：nums = [1,0,0,1,0,1], k = 2
输出：1
解释：在第一次操作时，nums 可以变成 [1,0,0,0,1,1] 得到连续两个 1 。
示例 2：
输入：nums = [1,0,0,0,0,0,1,1], k = 3
输出：5
解释：通过 5 次操作，最左边的 1 可以移到右边直到 nums 变为 [0,0,0,0,0,1,1,1] 。
示例 3：
输入：nums = [1,1,0,1], k = 2
输出：0
解释：nums 已经有连续 2 个 1 了。
提示：

1 <= nums.length <= 10⁵

nums[i] 要么是 0 ，要么是 1 。

1 <= k <= sum(nums)

解题思路#

对于 \(\texttt{nums}\) 中包含 \(k\) 个 \(1\) 的某个子序列，假设其下标为 \(a_0,\cdots,a_{k-1}\)，要将其变为下标连续的 \(b,b+1,\cdots,b+k-1\)，即要使得

\[ \sum_{i=0}^{k-1}=\vert a_i - (b + i) \vert=\sum_{i=0}^{k-1}\vert (a_i - i)-b \vert \]

最小。该问题就是 462. 最少移动次数使数组元素相等 II。

即当 \(b\) 为 \(a_i-i\) 的中位数时，和最小。

假设中位数的位置为 \(m=\lfloor \frac{k}{2}\rfloor\)，则中位数为 \(b=a[m]-m\)。

将 \(\{a_i\}\) 分成 \([a_0,a_{m-1}]\) 和 \([a_m,a_{k-1}]\) 两部分。

\[ \begin{aligned} \sum_{i=0}^{k-1}\vert (a_i - i)-b \vert &=\\ &=\left[m\times b-\left(\underbrace{\sum_{i=0}^{m-1}a_i}_{A_L}-\underbrace{\sum_{i=0}^{m-1}i}_{S_L}\right)\right]+\\ &\left[\left(\underbrace{\sum_{i=m}^{k-1}a_i}_{A_R}-\underbrace{\sum_{i=m}^{k-1}i}_{S_R}\right)-(k-m)\times b\right]\\ &=(2m-k)\times b+(A_R-A_L)+(S_L-S_R) \end{aligned} \]

from itertools import accumulate

class Solution:
    def minMoves(self, nums: List[int], k: int) -> int:
        if k == 1:
            return 0

        a = [i for i, x in enumerate(nums) if x == 1]
        tot = list(accumulate(a))

        M = k // 2
        SL = M * (M - 1) // 2 # [i,i+M-1]
        SR = (M + k - 1) * (k - M) // 2 # [i+M,i+k-1]
        S = SL - SR

        ans = float('inf')
        n = len(a)
        for i in range(n - k + 1):
            AL = tot[i + M - 1] - (i and tot[i - 1])
            AR = tot[i + k - 1] - tot[i + M - 1]
            b = a[i + M] - M # 中位数
            t = (2 * M - k) * b + AR - AL + S
            ans = min(ans, t)

        return ans

1703. 得到连续 K 个 1 的最少相邻交换次数#

问题描述#

解题思路#

1703. 得到连续 K 个 1 的最少相邻交换次数 #