使子数组元素和相等

标签: 数组 数学 数论 排序

难度: Medium

给你一个下标从 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

Submission

运行时间: 120 ms

内存: 33.2 MB

from itertools import accumulate
from math import gcd


class Solution:
    def makeSubKSumEqual(self, arr: list[int], k: int) -> int:

        n = len(arr)
        g = gcd(n, k)

        def calc_cost(r):
            sub_arr = arr[r::g]
            sub_len = len(sub_arr)
            sub_arr.sort()
            median = sub_arr[sub_len // 2]
            return sum(abs(x - median) for x in sub_arr)

        return sum(calc_cost(i) for i in range(g))

Explain

此题解采用了基于余数分组的策略,考虑了循环数组的特性和子数组总和的均等性。具体方法是通过计算数组长度n与子数组长度k的最大公约数g,将原数组根据余数分为g组,每组的元素是从数组索引r开始,每隔g个元素取一个形成的子数组。由于要求每个长度为k的子数组的和相同,这实际上转化为了确保在所有循环中跨越数组的这些固定偏移集合中的元素总和相同。对每个分组计算将其中所有元素调整为中位数的代价,然后累计这些代价以得到最终答案。选择中位数是因为中位数是最小化绝对偏差之和的最优策略,这在统计学中是一个众所周知的事实。

时间复杂度: O(n log(n/k))

空间复杂度: O(n)

from itertools import accumulate
from math import gcd

class Solution:
    def makeSubKSumEqual(self, arr: list[int], k: int) -> int:
        n = len(arr)
        g = gcd(n, k)  # 计算n和k的最大公约数

        def calc_cost(r):
            # 生成以r为起点,步长为g的子数组
            sub_arr = arr[r::g]
            sub_len = len(sub_arr)
            # 对子数组进行排序以便找到中位数
            sub_arr.sort()
            # 中位数
            median = sub_arr[sub_len // 2]
            # 计算调整到中位数的总代价
            return sum(abs(x - median) for x in sub_arr)

        # 计算所有分组的调整代价之和
        return sum(calc_cost(i) for i in range(g))

Explore

通过计算最大公约数g来分组的原因是,这样可以保证每组中的元素间隔g个元素,确保了这些元素在任何长度为k的子数组中均匀分布。由于长度为k的子数组需要覆盖整个原数组,且k可能不是n的因子,g作为n和k的最大公约数保证了每个子数组中的每个分组的元素可以均匀出现。这样,只要确保每个分组的元素总和相等,就可以保证所有长度为k的子数组的总和相等。

选择中位数作为目标值是因为中位数是统计学中一个重要概念,用于最小化一组数的绝对偏差之和。在数学理论中,这被称为最优化L1范数问题,即中位数是使得从一组数据到该点的绝对距离之和最小的点。这种性质使得在调整数组元素以使其和相等时,总的调整代价最小。

在给定的代码中,当数组长度为偶数时,选择的是下中位数(即中间位置的前一个元素)。这是因为在Python的列表切片中,使用整数除法(//)自动向下取整。选择上中位数或下中位数可能会在某些特定情况下影响总代价,但在大多数情况下,这种差异对于总体代价是微小的。选择下中位数简化了计算方法,且仍然能有效地减小总偏差。

算法中对递归深度的考虑基于gcd函数的递归实现。Euclid算法(最常用的求gcd的方法)的递归深度取决于两个数字的大小及其关系,最坏的情况发生在两个连续的Fibonacci数,此时递归深度可以达到O(log(min(a, b))),其中a和b是gcd函数的输入。因此在实际应用中,递归深度通常是可管理的,因为输入的数值范围通常有限。