可以借助np.compress
或np.where
实现,详见下边的代码示例:
import numpy as np
a_ndarray = np.array(
[
[1, 3],
[3, 4],
[3, 4444],
[1, 6],
[8, 6],
[10, 6]
]
)
# 按照第一列进行筛选,选中等于1的行
a_1 = a_ndarray[np.where(a_ndarray[:, 0] == 1)]
a_2 = np.compress(a_ndarray[:, 0] == 1, a_ndarray, axis=0)
# 检查发现两种方法得到的结果一致
check_1 = np.all(a_1 == a_2)
# 按照第一列进行筛选,选中等于1或3的行
a_3 = a_ndarray[
np.where(
(a_ndarray[:, 0] == 1) | (a_ndarray[:, 0] == 3)
)
]
a_4 = np.compress(
(a_ndarray[:, 0] == 1) | (a_ndarray[:, 0] == 3),
a_ndarray,
axis=0
)
# 检查发现两种方法得到的结果一致
check_2 = np.all(a_3 == a_4)
参考资料:
NumPy官方:numpy.compress