如何使用numpy原生方法高效按条件过滤二维numpy数组的行

最优实现方案 #

直接使用numpy原生布尔索引即可，全程在numpy底层运算，没有Python层面的循环和类型转换，性能远高于filter转列表的方案，代码更简洁：

result = my_arr[my_arr[:, 0] <= 1]

实现逻辑说明 #

• my_arr[:, 0] 提取二维数组所有行的第一列元素，得到一维数组
• my_arr[:, 0] <= 1 对第一列所有元素做向量化判断，生成和行数长度一致的布尔数组，满足条件的位置为True，反之为False
• 用布尔数组索引原二维数组，会自动保留布尔值为True对应的行，直接输出numpy数组类型的结果

示例验证 #

针对你给出的示例数组，运行效果如下：

import numpy as np
my_arr = np.array([[0, 4],
                   [1, 5],
                   [2, 6]])
result = my_arr[my_arr[:, 0] <= 1]
print(result)
# 输出：
# [[0 4]
#  [1 5]]

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