题目链接：2607. 使子数组元素和相等
方法：分组 + gcd + 中位数
解题思路
题意：将\(arr\)中某个元素\(+1\)或\(-1\)，使得任意长度为\(k\)的子数组的元素总和相等，且总操作数最少；
1、首先考虑数组\(arr\)为非循环数组：
- 任意\(k\)长的子数组总和相等，则有下述情形，依次可以将\(arr\)分组，使得组内的元素相同，此时对组内元素操作时，为了使得操作次数最小，应该使得每个组内元素与组内的中位数相同，此时所需次数最少。
/*
a[0] + a[1] + .. + a[k - 1]
= a[1] + a[2] + ... + a[k]
= a[2] + a[3] + ... + a[k + 1]
...
= a[n - k] + ... + a[n - 1]
=>
a[0] = a[k] = a[2k] = ... = a[xk], xk < n
a[1] = a[k + 1] = a[2k + 1] = ... = a[xk + 1], xk + 1 < n
a[2] = a[k + 2] = a[2k + 2] = ... = a[xk + 2], xk + 2 < n
...
*/
2、考虑数组\(arr\)为循环数组：
- 在上述的情况下，增加循环条件，每一组遍历到当前元素已经被添加过，因此需要有一个\(isvisit\)数组，判断当前元素是否已经被添加，然后使得每个组内元素与组内的中位数相同，输出总的操作次数。
/*
a[0] + a[1] + .. + a[k - 1]
= a[1] + a[2] + ... + a[k]
= a[2] + a[3] + ... + a[k + 1]
...
= a[n - k] + ... + a[n - 1]
= a[n - k + 1] + ... + a[n - 1] + a[0]
= a[n - k + 2] + ... + a[n - 1] + a[0] + a[1]
...
=>
a[0] = a[k] = a[2k] = ... = a[xk % n]
a[1] = a[k + 1] = a[2k + 1] = ... = a[(xk + 1) % n]
a[2] = a[k + 2] = a[2k + 2] = ... = a[(xk + 2) % n]
...
*/
代码
class Solution {
public:
long long makeSubKSumEqual(vector<int>& arr, int k) {
int n = arr.size();
long long ans = 0;
vector<int> isvisit(n);
for (int i = 0; i < k; i ++ ) {
if (isvisit[i]) continue; // 被添加过则跳过
vector<int> cur;
int j = i;
while (!isvisit[j]) { // isvisit[j] = 1，表明回到某一个被添加过的元素，终止遍历
isvisit[j] = 1;
cur.push_back(arr[j]);
j = (j + k) % n;
}
sort(cur.begin(), cur.end());
for (auto &num : cur) {
ans += abs(num - cur[cur.size() / 2]);
}
}
return ans;
}
};
复杂度分析
时间复杂度：\(O(nlogn)\)；
空间复杂度：\(O(n)\)。
优化
- 裴蜀定理简介
若\(a,b\)是整数，且\(gcd(a,b)=d\)，那么对于任意的整数\(x,y\)，\(ax+by\)都一定是\(d\)的倍数，特别地，一定存在整数\(x,y\)，使\(ax+by=d\)成立。
由于本题为循环数组，所以周期除了\(k\)还有\(n\)，那么对于\(a[i]\)有，\(a[i] = a[i + xk + yn]\)。根据上述定理，有\(a[i] = a[i + gcd(k, n)]\)，那么只需要判断。
/*
k = gdc(k, n)
a[0] = a[k] = a[2k] = ... = a[xk], xk < n
a[1] = a[k + 1] = a[2k + 1] = ... = a[xk + 1], xk + 1 < n
a[2] = a[k + 2] = a[2k + 2] = ... = a[xk + 2], xk + 2 < n
...
*/
class Solution {
public:
long long makeSubKSumEqual(vector<int>& arr, int k) {
int n = arr.size();
long long ans = 0;
k = gcd(n, k);
for (int i = 0; i < k; i ++ ) {
vector<int> cur;
for (int j = i; j < n; j += k) cur.push_back(arr[j]);
nth_element(cur.begin(), cur.begin() + cur.size() / 2, cur.end()); // 快速选择降低复杂度
for (auto &num : cur) {
ans += abs(num - cur[cur.size() / 2]);
}
}
return ans;
}
};
复杂度分析
时间复杂度：\(O(n)\)；
空间复杂度：\(O(n)\)。