如何用Numpy高效实现向量与另一向量各元素比较生成布尔矩阵

你可以直接利用Numpy的广播机制实现需求，完全避免Python层的循环，性能远高于列表推导式的实现，代码如下：

import numpy as np
a1 = np.array([1, 2, 3])
a2 = np.array([1, 2, 3])

# 核心实现：给a2增加列维度触发广播
result = a1 > a2[:, None]
# 等价写法：result = np.greater(a1, a2[:, np.newaxis])

print(result)

运行输出和你的原始代码完全一致：

[[False  True  True]
 [False False  True]
 [False False False]]

实现原理 #

• 原始的a2是形状为(3,)的一维向量，通过[:, None]（等价于[:, np.newaxis]）操作后，形状变为(3, 1)的二维数组
• 根据Numpy广播规则，形状为(3, 1)的数组和形状为(3,)的a1做比较运算时，两者会自动扩展为形状(3, 3)的数组后逐元素计算，最终得到的结果第i行刚好是a1和a2[i]逐元素比较的结果

性能优势 #

当向量长度很大时，该方案的执行速度是列表推导式实现的几十到上百倍：所有运算都在Numpy底层C语言层面执行，没有Python循环的额外开销，也符合Numpy的向量运算最佳实践。

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