python - 有没有更快的方法来解决以下问题？

Question

A 是 mn 矩阵
B 是一个 nn 矩阵

我想返回大小为 m*n 的矩阵 C，使得:

在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^2)

如果A[i][j] - B[j][k]总是 > 0 它可以很容易地改进为

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

但当出现 A[i][j] - B[j][k]< 0 的情况时，也可以改进？我认为一些分而治之的算法可能会有所帮助，但我对它们不熟悉。

Answer 1

对于每个 j，您可以对每列 B[j][:] 进行排序并计算累积和。

然后对于给定的A[i][j]，您可以找到大于A[的B[j][k]之和使用二分搜索在 O(log n) 时间内完成 i][j]。如果 B[j][:] 中有 x 个元素大于 A[i][j] 并且它们的总和为 S，那么C[i][j] = A[i][j] * x - S。

这为您提供了一个整体 O((m+n)n log n) 时间算法。