将数组查询的时间复杂度从O(n²)优化至O(n)的方案问询

问题说明 #

给定数组和查询数组：

array = [1,3,5,6,9]
queries = [5,9,3]

要计算把数组里所有元素调整成每个查询值需要的总操作次数（每次只能加1或减1）。举几个实际计算的例子：

• 查询值为5时，总操作次数：abs(1-5)+abs(3-5)+abs(5-5)+abs(6-5)+abs(9-5) = 4+2+0+1+4 = 11
• 查询值为9时，总操作次数：21
• 查询值为3时，总操作次数：13

现有一段Python实现代码，但时间复杂度为O(len(array) × len(queries))，数据量大时效率很低，需要更优的实现方案：

def find_min_costs(array,queries):
    for query in queries:
        cost = 0
        for arr in array:
            cost+= abs(arr-query)
        print(cost,end = ' ')

优化方案 #

核心思路 #

不用逐个元素计算差值，靠排序+前缀和+二分查找就能大幅提速。先给原数组排序，再计算前缀和数组，这样每个查询只需用二分法找到分界点，就能通过数学公式直接算出总操作次数，避免遍历整个数组。

具体步骤：

1. 给原数组排序，让元素从小到大排列，方便用二分法把数组分成「小于等于查询值」和「大于查询值」两部分。
2. 计算排序数组的前缀和数组，前缀和能快速求出任意前k个元素的总和，不用每次重复累加。
3. 对每个查询值，用二分法找到第一个比它大的元素的索引：
   • 索引左边的元素都≤查询值，这部分总操作次数 = 查询值×左边元素个数 - 左边元素总和
   • 索引右边的元素都>查询值，这部分总操作次数 = 右边元素总和 - 查询值×右边元素个数
   • 两部分结果相加就是当前查询对应的总操作次数

优化后的代码 #

import bisect

def find_min_costs_optimized(array, queries):
    sorted_arr = sorted(array)
    n = len(sorted_arr)
    # 前缀和数组：prefix_sum[0]=0，prefix_sum[1]是第一个元素的和，prefix_sum[n]是整个数组的总和
    prefix_sum = [0] * (n + 1)
    for i in range(n):
        prefix_sum[i+1] = prefix_sum[i] + sorted_arr[i]
    
    results = []
    for q in queries:
        # 找到第一个大于q的元素的位置
        k = bisect.bisect_right(sorted_arr, q)
        # 计算左边部分的操作次数
        left_cost = q * k - prefix_sum[k]
        # 计算右边部分的操作次数
        right_cost = (prefix_sum[n] - prefix_sum[k]) - q * (n - k)
        total_cost = left_cost + right_cost
        results.append(str(total_cost))
    
    print(' '.join(results))

# 测试示例
array = [1,3,5,6,9]
queries = [5,9,3]
find_min_costs_optimized(array, queries)  # 输出：11 21 13

时间复杂度说明 #

• 排序数组：O(n log n)
• 计算前缀和：O(n)
• 每个查询的二分查找+计算：O(log n)，m个查询总耗时O(m log n)
• 总体时间复杂度为O(n log n + m log n)，当数组长度和查询数量都较大时，远优于原方案的O(n*m)

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