跳转至

2607. 使子数组元素和相等

题目描述

给你一个下标从 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 <= 105
  • 1 <= arr[i] <= 109

解法

方法一:数学(裴蜀定理) + 中位数贪心

题目要求每个长度为 $k$ 的子数组的元素总和相等,那么有以下等式成立:

$$ arr_i + arr_{i + 1} + \cdots + arr_{i + k - 1} = arr_{i + 1} + arr_{i + 2} + \cdots + arr_{i + k} $$

化简得:

$$ arr_i = arr_{i + k} $$

也即是说,数组 $arr$ 有一个大小为 $k$ 的循环节,而由于数组 $arr$ 是一个循环数组,那么数组 $arr$ 也有一个长度为 $n$ 的循环节。换句话说,数组 $arr$ 上间隔为 $k$,以及间隔为 $n$ 的元素均相等。即有:

$$ arr_i = arr_{i + k \times x + n \times y} $$

根据裴蜀定理,有 $a \times x + b \times y = gcd(a, b)$,因此,有:

$$ arr_i = arr_{i + k \times x + n \times y} = arr_{i + gcd(k, n)} $$

因此,数组 $arr$ 上的元素可以分为 $gcd(k, n)$ 组,每组的元素间隔为 $gcd(k, n)$,且每一组中的所有元素均相等。对于每一组,我们可以将其元素按照大小排序,然后取中位数,即可将该组中的所有元素变为中位数。对于所有组,我们将其中位数之差的绝对值求和,即为答案。

时间复杂度 $O(n \times \log n)$,空间复杂度 $O(n)$。其中 $n$ 为数组 $arr$ 的长度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
    def makeSubKSumEqual(self, arr: List[int], k: int) -> int:
        n = len(arr)
        g = gcd(n, k)
        ans = 0
        for i in range(g):
            t = sorted(arr[i:n:g])
            mid = t[len(t) >> 1]
            ans += sum(abs(x - mid) for x in t)
        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
class Solution {
    public long makeSubKSumEqual(int[] arr, int k) {
        int n = arr.length;
        int g = gcd(n, k);
        long ans = 0;
        for (int i = 0; i < g; ++i) {
            List<Integer> t = new ArrayList<>();
            for (int j = i; j < n; j += g) {
                t.add(arr[j]);
            }
            t.sort((a, b) -> a - b);
            int mid = t.get(t.size() >> 1);
            for (int x : t) {
                ans += Math.abs(x - mid);
            }
        }
        return ans;
    }

    private int gcd(int a, int b) {
        return b == 0 ? a : gcd(b, a % b);
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
public:
    long long makeSubKSumEqual(vector<int>& arr, int k) {
        int n = arr.size();
        int g = gcd(n, k);
        long long ans = 0;
        for (int i = 0; i < g; ++i) {
            vector<int> t;
            for (int j = i; j < n; j += g) {
                t.push_back(arr[j]);
            }
            sort(t.begin(), t.end());
            int mid = t[t.size() / 2];
            for (int x : t) {
                ans += abs(x - mid);
            }
        }
        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
29
30
func makeSubKSumEqual(arr []int, k int) (ans int64) {
    n := len(arr)
    g := gcd(n, k)
    for i := 0; i < g; i++ {
        t := []int{}
        for j := i; j < n; j += g {
            t = append(t, arr[j])
        }
        sort.Ints(t)
        mid := t[len(t)/2]
        for _, x := range t {
            ans += int64(abs(x - mid))
        }
    }
    return
}

func abs(x int) int {
    if x < 0 {
        return -x
    }
    return x
}

func gcd(a, b int) int {
    if b == 0 {
        return a
    }
    return gcd(b, a%b)
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
function makeSubKSumEqual(arr: number[], k: number): number {
    const n = arr.length;
    const g = gcd(n, k);
    let ans = 0;
    for (let i = 0; i < g; ++i) {
        const t: number[] = [];
        for (let j = i; j < n; j += g) {
            t.push(arr[j]);
        }
        t.sort((a, b) => a - b);
        const mid = t[t.length >> 1];
        for (const x of t) {
            ans += Math.abs(x - mid);
        }
    }
    return ans;
}

function gcd(a: number, b: number): number {
    if (b === 0) {
        return a;
    }
    return gcd(b, a % b);
}

评论