如何在二维Numpy数组中基于指定列的取值筛选行？

需求说明 #

需要从以下Numpy数组中筛选出索引为1的列值为1且索引为2的列值为0的行：

import numpy as np

a = np.array([[1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 1, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0],
       [1, 0, 1, 1, 0]])

你的尝试问题分析 #

你写出的代码：

evidence = {1:1,2:0}
mask = a[:,list(evidence.keys())] == list(evidence.values())

得到的mask是二维布尔数组（每行对应原数组的一行，每个元素对应对应列的条件是否满足），但筛选行需要一维布尔数组（标记每行是否所有条件都满足），所以需要进一步处理。

完整解决代码 #

import numpy as np

a = np.array([[1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 1, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [1, 1, 1, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 1, 0, 0, 1],
       [1, 1, 0, 0, 0],
       [1, 0, 0, 1, 0],
       [1, 0, 1, 1, 0]])

evidence = {1:1,2:0}
cols = list(evidence.keys())
vals = list(evidence.values())
# 生成每行的条件判断结果，再按行取逻辑与，得到每行是否满足所有条件
mask = np.all(a[:, cols] == vals, axis=1)
# 用掩码筛选符合条件的行
result = a[mask]
print(result)

代码解释 #

1. a[:, cols] == vals：对比指定列的值是否等于目标值，得到每行的布尔判断结果
2. np.all(..., axis=1)：对每行的布尔值取逻辑与，只有当该行所有列都满足条件时，结果为True
3. a[mask]：用一维布尔掩码直接索引原数组，得到符合条件的所有行

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