题目描述
给你一个整数数组 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$ 的个数。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15 | 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
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28 | 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;
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29 | 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;
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24 | 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
}
|