NumPy如何不使用if语句仅提取条件匹配的数组元素

你当前使用的np.where(b =='yes',a,0)属于三目逻辑的逐元素映射：位置满足条件时取a对应值，不满足就取你传入的第三个参数0，所以返回结果和原数组长度一致，会保留不匹配位置的填充值，无法直接得到筛选后的压缩结果。

不需要写if条件判断，直接用NumPy原生布尔索引就能实现需求，代码如下：

import numpy as np

# 定义示例数组
a = np.array([27, 28, 29, 21, 17, 15, 19, 22, 18, 23, 24, 25, 30, 26])
b = np.array(['no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no', 'yes', 'yes', 'yes', 'yes', 'yes', 'no'], dtype=object)

# 布尔索引直接筛选
res = a[b == 'yes']

运行后res的输出完全匹配你的预期：

array([29, 21, 17, 19, 18, 23, 24, 25, 30])

实现原理 #

• b == 'yes'会生成一个和b长度相同的布尔型数组，匹配yes的位置值为True，其余位置为False
• 把这个布尔数组作为索引传入a时，NumPy会自动保留布尔值为True位置的a元素，丢弃False位置的元素，不会生成占位填充值，直接返回筛选后的元素集合

如果你习惯用np.where的写法，也可以取np.where返回的索引数组来做索引，效果完全一致：

res = a[np.where(b == 'yes')[0]]

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