如何根据列值筛选Numpy数组的行？

针对你需要筛选Numpy数组中某列值属于预定义列表的行的需求，可以用Numpy内置的np.isin()函数高效实现，以下是具体步骤：

示例代码 #

import numpy as np

# 原始数组
myArray = np.array([[1,55,4],
                    [2,2,3],
                    [3,90,2],
                    [4,65,1]])
# 目标值列表
desiredValues = [2,3,4]

# 1. 生成布尔掩码：判断第一列的每个元素是否在目标列表中
mask = np.isin(myArray[:, 0], desiredValues)
# 2. 用掩码筛选符合条件的行
desiredArray = myArray[mask]

print(desiredArray)

执行后输出结果：

[[ 2  2  3]
 [ 3 90  2]
 [ 4 65  1]]

关键说明 #

• myArray[:, 0]：提取数组的第一列（索引从0开始，若要筛选其他列，修改索引即可）。
• np.isin()：生成与第一列长度一致的布尔数组，元素为True表示对应位置的列值在目标列表中。
• 布尔掩码索引：直接用布尔数组作为Numpy数组的索引，会自动保留所有True对应的行。

备选方法（适合小数据量） #

如果数组规模很小，也可以用Python列表推导式实现，但效率远低于np.isin()的向量化操作：

desiredArray = myArray[[row[0] in desiredValues for row in myArray]]

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