如何优化代码执行时间与内存占用？附列表调整代码分析

问题背景 #

需求：降低代码执行时间与内存占用。
该代码需完成以下功能：给定s个列表，每个列表有n个数字，移除k个数字后将列表分为p1（左半部分）和p2（右半部分），目标是让sum(p1) - sum(p2)的差值最大化。

原代码的核心性能问题 #

• 重复复制列表：每次迭代都复制完整的p1和p2，内存开销随循环次数线性增长。
• 重复求和：每次操作后都重新计算sum(p1)和sum(p2)，时间复杂度高。
• 低效的极值操作：使用list.remove(min(p1))和list.remove(max(p2))，每次找极值和删除元素都是O(n)操作，累积开销大。
• 冗余存储：用字典存储所有可能的操作结果，浪费内存。

针对性优化方案 #

1. 用差值增量替代重复求和 #

不用每次修改列表后重新计算总和，而是基于初始差值计算操作带来的增量：

• 初始差值 diff = sum(p1) - sum(p2)
• 若将p2的元素x移到p1，差值增量为2*x（p1加x，p2减x，差值变化+2x）
• 若移除p1的元素x，差值增量为-x
• 若移除p2的元素x，差值增量为+x

2. 用堆结构快速获取极值 #

使用heapq模块维护p1的小顶堆和p2的大顶堆（用负号模拟），获取极值的时间复杂度从O(n)降到O(logn)，删除极值的操作也更高效。

3. 避免冗余存储与复制 #

不再存储所有可能的新列表，直接计算每个操作的差值增量，选择增量最大的操作执行，节省内存。

优化后的代码 #

import heapq

s = int(input())
answer = []
for _ in range(s):
    n, k = map(int, input().split())
    mobs_list = list(map(int, input().split()))
    mid = n // 2
    p1 = mobs_list[:mid]
    p2 = mobs_list[mid:]
    
    # 构建小顶堆（p1）和大顶堆（p2用负号模拟）
    heapq.heapify(p1)
    max_p2_heap = [-num for num in p2]
    heapq.heapify(max_p2_heap)
    
    sum_p1 = sum(p1)
    sum_p2 = sum(p2)
    remaining = k
    
    while remaining > 0:
        if remaining >= 2:
            gains = []
            # 操作1：移除p1末尾两个元素，将p2首元素移到p1
            if len(p1) >= 2 and len(p2) >= 1:
                gain = (-p1[-1] - p1[-2] + p2[0]) + p2[0]
                gains.append((gain, 1))
            # 操作2：移除p1末尾元素，移除p2最大值
            if len(p1) >= 1 and len(p2) >= 1:
                max_p2_val = -max_p2_heap[0]
                gain = -p1[-1] + max_p2_val
                gains.append((gain, 2))
            # 操作3：移除p1两个最小值，将p2首元素移到p1
            if len(p1) >= 2 and len(p2) >= 1:
                min1 = p1[0]
                temp = heapq.heappop(p1)
                min2 = p1[0]
                heapq.heappush(p1, temp)
                gain = (-min1 - min2 + p2[0]) + p2[0]
                gains.append((gain, 3))
            # 操作4：移除p2两个最大值，将p1末尾元素移到p2
            if len(p2) >= 2 and len(p1) >= 1:
                max1 = -max_p2_heap[0]
                temp = heapq.heappop(max_p2_heap)
                max2 = -max_p2_heap[0]
                heapq.heappush(max_p2_heap, temp)
                gain = max1 + max2
                gains.append((gain, 4))
            # 操作5：移除p1最小值，移除p2最大值
            if len(p1) >= 1 and len(p2) >= 1:
                min_p1 = p1[0]
                max_p2_val = -max_p2_heap[0]
                gain = -min_p1 + max_p2_val
                gains.append((gain, 5))
            
            # 选择增益最大的操作执行
            max_gain, op = max(gains, key=lambda x: x[0])
            if op == 1:
                sum_p1 -= p1[-1] + p1[-2]
                sum_p1 += p2[0]
                sum_p2 -= p2[0]
                del p1[-2:]
                del p2[0]
            elif op == 2:
                sum_p1 -= p1[-1]
                max_val = -heapq.heappop(max_p2_heap)
                sum_p2 -= max_val
                del p1[-1]
            elif op == 3:
                sum_p1 -= heapq.heappop(p1) + heapq.heappop(p1)
                sum_p1 += p2[0]
                sum_p2 -= p2[0]
                del p2[0]
            elif op == 4:
                sum_p2 -= (-heapq.heappop(max_p2_heap)) + (-heapq.heappop(max_p2_heap))
                sum_p2 += p1[-1]
                sum_p1 -= p1[-1]
                del p1[-1]
            elif op == 5:
                sum_p1 -= heapq.heappop(p1)
                sum_p2 -= (-heapq.heappop(max_p2_heap))
            remaining -= 2
        else:
            gains = []
            # 操作1：移除p1末尾元素，将p2首元素移到p1
            if len(p1) >= 1 and len(p2) >= 1:
                gain = (-p1[-1] + p2[0]) + p2[0]
                gains.append((gain, 1))
            # 操作2：移除p2首元素
            if len(p2) >= 1:
                gain = p2[0]
                gains.append((gain, 2))
            # 操作3：移除p1最小值，将p2首元素移到p1
            if len(p1) >= 1 and len(p2) >= 1:
                min_p1 = p1[0]
                gain = (-min_p1 + p2[0]) + p2[0]
                gains.append((gain, 3))
            # 操作4：移除p2最大值
            if len(p2) >= 1:
                max_p2_val = -max_p2_heap[0]
                gain = max_p2_val
                gains.append((gain, 4))
            
            max_gain, op = max(gains, key=lambda x: x[0])
            if op == 1:
                sum_p1 -= p1[-1]
                sum_p1 += p2[0]
                sum_p2 -= p2[0]
                del p1[-1]
                del p2[0]
            elif op == 2:
                sum_p2 -= p2[0]
                del p2[0]
            elif op == 3:
                sum_p1 -= heapq.heappop(p1)
                sum_p1 += p2[0]
                sum_p2 -= p2[0]
                del p2[0]
            elif op == 4:
                sum_p2 -= (-heapq.heappop(max_p2_heap))
            remaining -= 1
    
    answer.append(sum_p1 - sum_p2)

for item in answer:
    print(item)

优化效果 #

• 内存占用：从O(n*k)降至O(n)，不再重复复制列表和存储冗余操作结果。
• 执行时间：将多次O(n)操作替换为O(logn)的堆操作，避免重复求和，整体时间复杂度从O(kn)优化为O(klogn)，处理大规模数据时性能提升明显。

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