基于另一Numpy数组条件统计目标值出现次数的高效实现方法

嘿，这个需求用Numpy的向量化特性可以非常简洁高效地搞定，完全不需要复杂的循环或者冗余代码！

核心实现代码 #

直接一行就能得到结果：

import numpy as np

a = np.array([1,-1,1,1,-1,-1])
b = np.array([.75,.35,.7,.8,.2,.6])

# 统计b > 0.5时a中1的出现次数
c = np.sum((b > 0.5) & (a == 1))
print(c)  # 输出：3

原理解释 #

这里利用了Numpy的向量化布尔运算，步骤拆解一下：

1. b > 0.5 生成一个布尔数组，标记所有满足b条件的位置：array([ True, False, True, True, False, True])
2. a == 1 生成另一个布尔数组，标记a中等于目标值的位置：array([ True, False, True, True, False, False])
3. 用&做按位与运算，得到同时满足两个条件的位置：array([ True, False, True, True, False, False])
4. np.sum()会把布尔数组里的True（等价于1）加起来，就是我们要的统计次数

灵活扩展方案 #

如果需要频繁更换b的条件或者a的目标值，可以封装成一个通用函数：

def count_conditional(a, b, condition_b, target_value):
    return np.sum(condition_b(b) & (a == target_value))

# 调用示例：比如统计b < 0.4时a中-1的次数
result = count_conditional(a, b, lambda x: x < 0.4, -1)
print(result)  # 输出：2

效率优势 #

这种方法比Python原生循环快得多——因为Numpy的向量化操作是底层用C实现的，避开了Python解释器的循环开销，当数组规模越大（比如百万级元素），性能提升越明显。

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