Leetcode.2607 使子数组元素和相等

题目链接

Leetcode.2607 使子数组元素和相等 Rating ： 2071

题目描述

给你一个下标从 0 开始的整数数组 arr和一个整数 k 。数组 arr是一个循环数组。换句话说，数组中的最后一个元素的下一个元素是数组中的第一个元素，数组中第一个元素的前一个元素是数组中的最后一个元素。

你可以执行下述运算任意次：

• 选中 arr中任意一个元素，并使其值加上 1 或减去 1 。
  执行运算使每个长度为 k 的 子数组 的元素总和都相等，返回所需要的最少运算次数。

子数组 是数组的一个连续部分。

示例 1：

输入：arr = [1,4,1,3], k = 2
输出：1
解释：在下标为 1 的元素那里执行一次运算，使其等于 3 。
执行运算后，数组变为 [1,3,1,3] 。

• 0 处起始的子数组为 [1, 3] ，元素总和为 4
• 1 处起始的子数组为 [3, 1] ，元素总和为 4
• 2 处起始的子数组为 [1, 3] ，元素总和为 4
• 3 处起始的子数组为 [3, 1] ，元素总和为 4

示例 2：

输入：arr = [2,5,5,7], k = 3
输出：5
解释：在下标为 0 的元素那里执行三次运算，使其等于 5 。在下标为 3 的元素那里执行两次运算，使其等于 5 。
执行运算后，数组变为 [5,5,5,5] 。

• 0 处起始的子数组为 [5, 5, 5] ，元素总和为 15
• 1 处起始的子数组为 [5, 5, 5] ，元素总和为 15
• 2 处起始的子数组为 [5, 5, 5] ，元素总和为 15
• 3 处起始的子数组为 [5, 5, 5] ，元素总和为 15

提示：

• $1<=k<=arr.length<=10^5$
• $1<=arr[i]<=10^9$

解法：裴蜀定理 + 中位数贪心

如题 $arr$ 是一个元素个数为 $n$ 的 循环数组，即它的周期为 $n$。

题目要求 $arr$ ，长度为 $k$ 的子数组和相等，即：

$$a_i+a_{i+1}+a_{i+2}+...+a_{i+k-1}=a_{i+1}+a_{i+2}+a_{i+3}+...+a_{i+k-1}+a_{i+k}$$

将两边化简得：

$$a_i=a_{i+k}$$

故 $k$ 也是 $arr$的周期。

根据结论：一个循环数组如果既有周期 $n$ ，又有周期 $k$ ，那么一定有周期 $g=gcd(n,k)$。

所以我们可以将原数组 $arr$ 的下标$i$，按 $i\%g$ 分到若干组内。对于每一组，我们分别求让整个组元素都相同的最小操作次数。

要求最小次数，我们需要先对每一个组排序，让组内的每个元素，减去组内的中位数，即是最小操作次数。

这里我们可以使用 快速排序 ，快速求中位数。

时间复杂度： $O(nlogn)$

C++代码：

using LL = long long;

class Solution {
public:
    long long makeSubKSumEqual(vector<int>& arr, int k) {
        int n = arr.size();

        int g = gcd(n,k);
        vector<vector<int>> a(n);
        for(int i = 0;i < n;i++){
            a[i % g].push_back(arr[i]);
        };

        LL ans = 0;
        for(auto &b:a){
            int len = b.size();
            // 快速排序求中位数
            nth_element(b.begin(),b.begin() + len / 2,b.end());
            for(auto x:b){
                ans += abs(x - b[len/2]);
            }
        }

        return ans;
    }
};