在二维数组中搜索多个元素？

Question

我有一个一维数组 match 和一个二维数组 data 。我可以像这样搜索一个或多个元素：

 data[data[:,0] == match[0] ]
 data[data[:,0] == match[0] & data[:,0] == match[1] ]

但是你如何搜索所有这些:

 data[ data[:,0] == match ]

---

我想我可以在 in1d() 中使用 numba 函数！！ :

    mask = np.zeros(len(ar1), dtype=bool)
    for a in ar2:
            mask |= (ar1 == a)
    return mask

Answer 1

您可以使用 numpy.isin:

np.isin(data[:, 0], match)

---

示例：

data = np.array([[1, 4],[5, 2],[2, 4]])
match = np.array([2, 4])

np.isin(data[:,0], match)
# array([False, False,  True])

data[np.isin(data[:, 0], match)]
# array([[2, 4]])