Numpy中外层最小向量化的后续优化问题

核心思路 #

要避免原方案的行循环，同时不生成超大张量，我们可以逐列处理：对矩阵的每一列，计算该列所有元素对的最小值构成的N×N矩阵，再将所有列的结果累加，最终得到目标矩阵。这种方式的内存峰值仅为O(N²)（存储结果矩阵）+ O(N)（存储当前列），完全符合内存限制要求。

实现代码 #

基础Numpy版本

import numpy as np

A = np.array([[1, 2], [3, 4], [5, 6]])
n_rows = A.shape[0]
output = np.zeros((n_rows, n_rows), dtype=A.dtype)

# 逐列处理，累加每列的元素对最小值矩阵
for col in A.T:
    output += np.minimum.outer(col, col)

print(output)
# 输出：
# [[ 3  3  3]
#  [ 3  7  7]
#  [ 3  7 11]]

加速版本（Numba JIT）

如果矩阵规模较大（M或N数值较高），可以用Numba将循环编译为机器码，进一步提升速度：

import numpy as np
from numba import njit

@njit
def compute_min_sum(A):
    n, m = A.shape
    output = np.zeros((n, n), dtype=A.dtype)
    for k in range(m):
        col = A[:, k]
        for i in range(n):
            ci = col[i]
            for j in range(n):
                output[i, j] += min(ci, col[j])
    return output

A = np.array([[1, 2], [3, 4], [5, 6]])
output = compute_min_sum(A)
print(output)

原理说明 #

对于原问题中第i行第j列的结果，本质是对所有列m取min(A[i,m], A[j,m])后求和。逐列处理时，每列生成的np.minimum.outer(col, col)恰好是该列所有(i,j)对的最小值矩阵，累加所有列的矩阵就等价于对每个(i,j)对的所有列最小值求和，完全匹配需求。

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