NumPy：按np.array过滤行

Question

我想通过检查另一个数组是否包含列值来过滤NumPy 2-d数组。我怎么能这样做？

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]])

Answer 2

您可以使用np.in1d检查in的第一列中的哪些值也位于ar中，然后使用返回的布尔索引，而不是使用another_ar获取ar：

的行

>>> ar[np.in1d(ar[:,0], another_ar)]
array([[  1,   2],
       [  6, -15]])

这可能比使用任何类型的for循环并使用in测试成员资格要快得多。