python - NumPy:通过 np.array 过滤行

Question

我想通过检查另一个数组是否包含列值来过滤 NumPy 二维数组。我该怎么做？

import numpy as np

ar = np.array([[1,2],[3,-5],[6,-15],[10,7]])
another_ar = np.array([1,6])
new_ar = ar[ar[:,0] in another_ar]
print new_ar

我希望得到 [[1,2],[6,-15]] 但上面的代码只打印 [1,2]。

Answer 1

您可以使用 np.where ，但请注意，由于 ar[:,0] 是第一个元素的列表，如果 ar 您需要遍历它并检查成员资格:

>>> ar[np.where([i in another_ar for i in ar[:,0]])]
array([[  1,   2],
       [  6, -15]])