Question

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

English Version

题目描述

给你一个整数数组 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 <= 105
• nums[i] 要么是 0 ，要么是 1 。
• 1 <= k <= sum(nums)

解法

方法一：前缀和 + 中位数枚举

我们可以将数组 $nums$ 中的 $1$ 的下标存入数组 $arr$ 中。接下来，我们预处理数组 $arr$ 的前缀和数组 $s$，其中 $s[i]$ 表示数组 $arr$ 中前 $i$ 个元素的和。

对于长度为 $k$ 的子数组，左侧（包含中位数）的元素个数 $x=\frac{k+1}{2}$，右侧的元素个数为 $y=k-x$。

我们枚举中位数的下标 $i$，其中 $x-1\leq i\leq len(arr)-y$，那么左侧数组的前缀和 $ls=s[i+1]-s[i+1-x]$，右侧数组的前缀和 $rs=s[i+1+y]-s[i+1]$。当前中位数在 $nums$ 中的下标为 $j=arr[i]$，将左侧 $x$ 个元素移动到 $[j-x+1,..j]$ 所需要的操作次数为 $a=(j+j-x+1)\times\frac{x}{2}-ls$，将右侧 $y$ 个元素移动到 $[j+1,..j+y]$ 所需要的操作次数为 $b=rs-(j+1+j+y)\times\frac{y}{2}$，那么总的操作次数为 $a+b$，我们取所有总的操作次数的最小值即可。

时间复杂度 $O(n)$，空间复杂度 $O(m)$。其中 $n$ 和 $m$ 分别为数组 $nums$ 的长度以及数组 $nums$ 中 $1$ 的个数。

class Solution:
  def minMoves(self, nums: List[int], k: int) -> int:
    arr = [i for i, x in enumerate(nums) if x]
    s = list(accumulate(arr, initial=0))
    ans = inf
    x = (k + 1) // 2
    y = k - x
    for i in range(x - 1, len(arr) - y):
      j = arr[i]
      ls = s[i + 1] - s[i + 1 - x]
      rs = s[i + 1 + y] - s[i + 1]
      a = (j + j - x + 1) * x // 2 - ls
      b = rs - (j + 1 + j + y) * y // 2
      ans = min(ans, a + b)
    return ans

class Solution {
  public int minMoves(int[] nums, int k) {
    List<Integer> arr = new ArrayList<>();
    int n = nums.length;
    for (int i = 0; i < n; ++i) {
      if (nums[i] != 0) {
        arr.add(i);
      }
    }
    int m = arr.size();
    int[] s = new int[m + 1];
    for (int i = 0; i < m; ++i) {
      s[i + 1] = s[i] + arr.get(i);
    }
    long ans = 1 << 60;
    int x = (k + 1) / 2;
    int y = k - x;
    for (int i = x - 1; i < m - y; ++i) {
      int j = arr.get(i);
      int ls = s[i + 1] - s[i + 1 - x];
      int rs = s[i + 1 + y] - s[i + 1];
      long a = (j + j - x + 1L) * x / 2 - ls;
      long b = rs - (j + 1L + j + y) * y / 2;
      ans = Math.min(ans, a + b);
    }
    return (int) ans;
  }
}

class Solution {
public:
  int minMoves(vector<int>& nums, int k) {
    vector<int> arr;
    for (int i = 0; i < nums.size(); ++i) {
      if (nums[i]) {
        arr.push_back(i);
      }
    }
    int m = arr.size();
    long s[m + 1];
    s[0] = 1;
    for (int i = 0; i < m; ++i) {
      s[i + 1] = s[i] + arr[i];
    }
    long ans = 1L << 60;
    int x = (k + 1) / 2;
    int y = k - x;
    for (int i = x - 1; i < m - y; ++i) {
      int j = arr[i];
      int ls = s[i + 1] - s[i + 1 - x];
      int rs = s[i + 1 + y] - s[i + 1];
      long a = (j + j - x + 1L) * x / 2 - ls;
      long b = rs - (j + 1L + j + y) * y / 2;
      ans = min(ans, a + b);
    }
    return ans;
  }
};

func minMoves(nums []int, k int) int {
  arr := []int{}
  for i, x := range nums {
    if x != 0 {
      arr = append(arr, i)
    }
  }
  s := make([]int, len(arr)+1)
  for i, x := range arr {
    s[i+1] = s[i] + x
  }
  ans := 1 << 60
  x := (k + 1) / 2
  y := k - x
  for i := x - 1; i < len(arr)-y; i++ {
    j := arr[i]
    ls := s[i+1] - s[i+1-x]
    rs := s[i+1+y] - s[i+1]
    a := (j+j-x+1)*x/2 - ls
    b := rs - (j+1+j+y)*y/2
    ans = min(ans, a+b)
  }
  return ans
}