优化计算f(i)=min dist(i,S)函数的Python代码执行效率

原代码核心性能瓶颈 #

• 每次查找元素下标时调用sorted_arr.index(i)和indexed.count(i)，单次操作时间复杂度为O(n)，n次循环后整体复杂度达到O(n²)，数据量稍大就会出现明显卡顿
• 每次循环内对距离数组做排序操作，单次排序复杂度为O(k log k)，累计复杂度为O(nk log k)，属于不必要的性能开销
• 存在大量数组切片、拼接、重复创建新容器的操作，进一步拖慢运行速度

优化思路 #

我们可以利用排序数组的单调性+前缀和数组，把整体时间复杂度降到O(n log n)，大幅提升运行效率：

1. 预处理阶段给每个元素绑定原始下标后再排序，解决重复元素下标定位的问题，避免O(n)的查找操作
2. 提前构建排序数组的前缀和数组，对于任意位置的元素，计算k个最近元素的差分之和可以直接用前缀和O(1)算出，不需要逐个计算累加
3. 利用排序数组的单调性，和当前元素差值最小的k个元素一定是其左右相邻的连续区间，直接用双指针定位窗口即可，不需要对距离做排序

优化后代码 #

n, k = map(int, input().split())
arr = list(map(int, input().split()))
# 绑定原下标排序，解决重复元素定位问题
sorted_with_idx = sorted((val, i) for i, val in enumerate(arr))
sorted_vals = [x[0] for x in sorted_with_idx]
# 记录原数组每个元素在排序数组中的位置
pos_map = [0] * n
for pos in range(n):
    original_idx = sorted_with_idx[pos][1]
    pos_map[original_idx] = pos

# 预处理前缀和数组
prefix = [0] * (n + 1)
for i in range(n):
    prefix[i+1] = prefix[i] + sorted_vals[i]

res = []
for i in range(n):
    pos = pos_map[i]
    x = sorted_vals[pos]
    # 双指针找k个最近的元素窗口
    left = max(0, pos - k)
    right = min(n-1, pos + k)
    # 调整窗口大小为k，排除当前元素自身
    cnt = right - left
    while cnt > k:
        if x - sorted_vals[left] > sorted_vals[right] - x:
            left += 1
        else:
            right -= 1
        cnt -= 1
    # 用前缀和计算差的和
    total = 0
    # 左边部分的和：x * 左边个数 - 左边区间和
    left_cnt = max(0, pos - left)
    if left_cnt > 0:
        total += x * left_cnt - (prefix[pos] - prefix[left])
    # 右边部分的和：右边区间和 - x * 右边个数
    right_cnt = right - pos
    if right_cnt > 0:
        total += (prefix[right+1] - prefix[pos+1]) - x * right_cnt
    res.append(str(total))

print(' '.join(res))

内容的提问来源于stack exchange，提问作者Михаил Лепин