如何优化大数据量下数组统一为指定数的最少操作数计算？

问题分析 #

原代码answer = [sum(abs(x-y) for x in a1) for y in a2]的时间复杂度为O(n*m)，当a1和a2的长度都达到10^{5量级时，总操作次数会飙升至10}10，完全无法满足效率要求。

优化方案：前缀和+二分查找 #

通过排序+前缀和数组，结合二分查找，可将时间复杂度降至O(n log n + m log n)，足以处理大规模数据。具体步骤如下：

• 对a1进行排序，为二分查找和前缀和计算做准备
• 计算a1的前缀和数组，快速获取任意前k个元素的总和
• 对a2中的每个值y，用二分查找确定a1中小于等于y的元素数量，再通过前缀和快速计算总操作数

代码实现 #

import bisect

def calculate_total_operations(a1, a2):
    a1_sorted = sorted(a1)
    n = len(a1_sorted)
    # 前缀和数组：prefix[0] = 0，prefix[k] = sum(a1_sorted[:k])
    prefix = [0] * (n + 1)
    for i in range(n):
        prefix[i+1] = prefix[i] + a1_sorted[i]
    
    answer = []
    for y in a2:
        # 找到第一个大于y的元素索引，即小于等于y的元素个数为k
        k = bisect.bisect_right(a1_sorted, y)
        # 计算两部分操作数并求和
        left_ops = y * k - prefix[k]
        right_ops = (prefix[n] - prefix[k]) - y * (n - k)
        answer.append(left_ops + right_ops)
    return answer

逻辑说明 #

1. 排序：将a1排序后，小于等于y的元素集中在左侧，大于y的在右侧，方便拆分计算。
2. 前缀和数组：prefix[k]存储a1_sorted前k个元素的总和，能在O(1)时间内获取任意区间的元素和。
3. 二分查找：用bisect.bisect_right快速定位y在排序数组中的插入点，得到小于等于y的元素数量k，时间复杂度O(log n)。
4. 总操作数拆分：
   • 左侧k个元素到y的总步数：每个元素需增加y-x，总和为y*k - sum(x)
   • 右侧n-k个元素到y的总步数：每个元素需减少x-y，总和为sum(x) - y*(n-k)
   • 两部分相加即为最终总操作数

内容的提问来源于stack exchange，提问作者user3525805