如何基于运算符o比较NumPy数组val与arr中的值？

问题背景 #

给定两个NumPy数组val和arr，以及一个运算符元组o，需要对val的每一行val_i执行以下操作：

• 将val_i的第j个元素与arr的第j行所有元素，通过运算符o[j]进行比较，得到布尔向量
• 将所有布尔向量执行按位与操作，最终得到每个val_i对应的结果向量

示例输入与预期输出如下：

import numpy as np
from operator import gt, lt

val = np.array([[3,7,1], [4,8,5], [5,10,3]])
arr = np.array([[1,2,3,4,5], [6,7,8,9,10], [9,7,5,3,1]])
o = (gt, gt, lt)

# 预期结果
result = np.array([[True, False, False, False, False],
                   [True, True, False, False, False],
                   [True, True, True, False, False]])

原实现采用嵌套循环，代码繁琐，需要更简洁的实现方式。

简洁实现方案 #

利用NumPy的广播机制和向量化操作，可以完全避免循环，实现高效简洁的代码：

import numpy as np
from operator import gt, lt

val = np.array([[3,7,1], [4,8,5], [5,10,3]])
arr = np.array([[1,2,3,4,5], [6,7,8,9,10], [9,7,5,3,1]])
o = (gt, gt, lt)

# 生成每个维度的比较结果：每个元素是形状为(3,5)的布尔数组
comparison_list = [op(val[:, j], arr[j]) for j, op in enumerate(o)]
# 对所有比较结果按元素执行按位与操作
result = np.logical_and.reduce(comparison_list)

print(result)
# 输出与预期一致：
# [[ True False False False False]
#  [ True  True False False False]
#  [ True  True  True False False]]

代码解释 #

1. 广播比较：
   • val[:, j]提取val的第j列（形状为(3,)），arr[j]提取arr的第j行（形状为(5,)）
   • NumPy会自动将两个数组广播为(3,5)的形状，完成每个val行的第j元素与arr第j行所有元素的比较，得到对应布尔数组
2. 按位与合并：
   • np.logical_and.reduce会遍历comparison_list中的所有布尔数组，对每个位置的元素依次执行按位与操作，最终合并为一个(3,5)的结果数组

这种方式完全利用NumPy的向量化特性，代码简洁且性能远优于循环实现，尤其适合处理大规模数组。

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