如何高效计算所有i,j对应max(Ai+Bj, Bi+Aj)的双重求和值

首先我们先对原式的max项做等价变形，简化计算逻辑：
max(Ai + Bj, Bi + Aj) 的大小关系等价于比较 Ai - Bi 和 Aj - Bj：

• 当 Ai - Bi >= Aj - Bj 时，max(Ai + Bj, Bi + Aj) = Ai + Bj
• 当 Ai - Bi < Aj - Bj 时，max(Ai + Bj, Bi + Aj) = Bi + Aj

优化思路 #

我们定义 Ci = Ai - Bi，对所有元素按Ci升序排序后，就可以通过前缀和、后缀和快速统计每个i对应的所有j的求和结果，避免两层循环：

1. 构造元素组 (Ci, Ai, Bi)，按Ci从小到大排序，时间复杂度O(N log N)
2. 预处理排序后B数组的前缀和数组prefix_b，prefix_b[k]表示前k个元素的B值之和，时间复杂度O(N)
3. 预处理排序后A数组的后缀和数组suffix_a，suffix_a[k]表示第k个元素之后所有元素的A值之和，时间复杂度O(N)
4. 遍历排序后的每个元素，按下面的公式累加总结果：
   对于排序后的第k个元素（从0开始计数），有k+1个j满足Cj <= Ck，剩余N - k -1个j满足Cj > Ck，对应求和项为：
   (k+1)*Ak + prefix_b[k+1] + (N -k -1)*Bk + suffix_a[k+1]
   遍历累加的时间复杂度为O(N)

整体时间复杂度为O(N log N)，远优于暴力算法的O(N²)，适合N较大的场景。 #

代码实现（Python示例） #

def calculate_z(A: list[int], B: list[int]) -> int:
    n = len(A)
    # 按Ci = Ai - Bi 升序排序
    sorted_arr = sorted((A[i] - B[i], A[i], B[i]) for i in range(n))
    # 预处理B的前缀和
    prefix_b = [0] * (n + 1)
    for idx in range(n):
        prefix_b[idx + 1] = prefix_b[idx] + sorted_arr[idx][2]
    # 预处理A的后缀和
    suffix_a = [0] * (n + 1)
    for idx in range(n-1, -1, -1):
        suffix_a[idx] = suffix_a[idx + 1] + sorted_arr[idx][1]
    total = 0
    for k in range(n):
        cnt_le = k + 1
        a_k = sorted_arr[k][1]
        b_k = sorted_arr[k][2]
        total += cnt_le * a_k + prefix_b[cnt_le] + (n - cnt_le) * b_k + suffix_a[k + 1]
    return total

正确性验证 #

举个简单测试用例：
输入A = [1,3]，B = [2,4]
暴力计算结果：max(3,3)+max(5,5)+max(5,5)+max(7,7) = 3+5+5+7=20
调用上述函数返回结果为20，和暴力计算结果一致。

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