如何用NumPy向量化操作计算矩阵向量对的逐元素最小值之和

给定维度为(n, m)的矩阵M，我们需要生成(n, n)的矩阵K，其中K[i,j]是M第i行与第j行逐元素取最小值后的向量总和。直接循环处理大规模矩阵效率极低，以下是纯NumPy向量化的解决方案：

核心思路 #

利用NumPy的广播机制实现行与行之间的逐元素比较，避免不必要的高维度数组生成：

1. 将M的行分别扩展维度：M[:, None, :]（形状(n,1,m)）和M[None, :, :]（形状(1,n,m)）
2. 对两个扩展后的数组逐元素取最小值，得到形状为(n,n,m)的中间数组（每个(i,j,:)对应第i行和第j行的逐元素最小值向量）
3. 沿第三个维度（axis=2）求和，得到最终的(n,n)矩阵K

代码实现 #

import numpy as np

# 示例矩阵
M = np.array([[2, 5, 4], [3, 1, 6]])
# 向量化计算K
K = np.minimum(M[:, None, :], M[None, :, :]).sum(axis=2)

print(K)
# 输出结果：
# [[11  7]
#  [ 7 10]]

为什么不用np.minimum.outer？ #

np.minimum.outer(M)会生成形状为(n,m,n,m)的数组，这是元素与元素之间的外操作，而非行与行的逐元素操作，不仅维度冗余、占用大量内存，后续求和也无法直接得到目标结果，完全不符合需求。而广播的方法仅生成(n,n,m)的中间数组，计算效率和内存利用率都远高于outer方法。

性能优势 #

对于大规模矩阵，这种向量化操作的效率比Python循环高10~1000倍（具体取决于矩阵规模），因为NumPy的底层运算由C实现，避免了Python解释器的循环开销。

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