如何优化满足特定求和规则的矩阵C的计算效率？原算法复杂度O(m*n²)

已知A是m×n矩阵，B是n×n矩阵，需生成大小为m×n的矩阵C，其中C的元素满足公式：
$C_{ij} = \sum_{k=1}^{n} max(0, a_{ij} - b_{jk}) $

当前Python实现代码如下：

for i in range(m):
   for j in range(n):
      C[i][j] = 0
      for k in range(n):
         C[i][j] += max(0, A[i][j] - B[j][k])

该算法时间复杂度为O(m*n²)。当A[i][j] - B[j][k]恒大于0时，可优化为：

C[i][j] = n*A[i][j] - sum(B[j]) 

但当存在A[i][j] - B[j][k] < 0的情况时，是否仍可优化？猜测分治算法可能有用，但并不熟悉这类算法。

当然可以优化，核心是通过预处理B的行，避免重复计算每个max(0, a_ij - b_jk)的累加，具体如下：

1. 预处理B的每一行 #

对B的每一行j执行两步操作：

• 将该行元素升序排序，得到B_sorted[j]
• 计算排序后数组的前缀和数组prefix_sum[j]，其中prefix_sum[j][t]表示前t个元素的总和（prefix_sum[j][0] = 0，prefix_sum[j][1] = B_sorted[j][0]，以此类推）

预处理的时间复杂度是O(n² log n)，仅需执行一次。

2. 快速计算C的元素 #

对于每个C[i][j]，我们只需要找到B的第j行中所有小于等于A[i][j]的元素：

• 用二分查找找到第一个大于A[i][j]的元素索引t，这意味着前t个元素都满足b_jk ≤ A[i][j]
• 这些元素对应的max(0, a_ij - b_jk)之和可以直接用公式计算：t*A[i][j] - prefix_sum[j][t]，剩下的元素因为b_jk > A[i][j]，贡献为0，无需计算

每个C[i][j]的计算时间为O(log n)，整体时间复杂度变为O(n² log n + m*n log n)。当m、n较大时，这个复杂度远低于原算法的O(mn²)（比如n=1000时，logn≈10，mnlogn是原复杂度的1/100）。

3. 关于分治算法的说明 #

分治确实能解决这类区间求和问题，比如构建线段树或归并树来支持快速查询符合条件的元素数量和总和，但这类方法实现复杂度更高。上面的排序+二分+前缀和方案实现简单，效率也足够满足大部分场景需求。

优化后代码示例 #

import bisect

# 预处理B的每一行
B_sorted = []
prefix_sum = []
n = len(B)
for j in range(n):
    sorted_row = sorted(B[j])
    B_sorted.append(sorted_row)
    # 计算前缀和
    ps = [0] * (n + 1)
    for idx in range(n):
        ps[idx + 1] = ps[idx] + sorted_row[idx]
    prefix_sum.append(ps)

# 计算C矩阵
m = len(A)
C = [[0] * n for _ in range(m)]
for i in range(m):
    for j in range(n):
        a_val = A[i][j]
        # 找到第一个大于a_val的元素位置
        count = bisect.bisect_right(B_sorted[j], a_val)
        C[i][j] = count * a_val - prefix_sum[j][count]

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