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

Question

我有一个 1d array match和 2d array 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

标签： numpysearchfilternumba

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