python - 如何获取包含在另一个数组中的一个数组的值的所有索引？

Question

假设我们有两个简单的一维 numpy 数组:

a = np.array([1, 1, 2, 5, 7, 8, 2, 4, 5, 6]) 
b = np.array([1, 5, 7])

现在我想要的是获取所有可能的索引，其中数组 b 中的每个值都包含在数组 a 中。

我们可以做，例如以下内容:

idx = np.where(np.any(a.reshape((-1, 1))==b, axis=1))[0]

idx 是 array([0, 1, 3, 4, 8]) (这是我真正想要的)。

现在我真的很好奇在 numpy 或任何其他库中是否已经有针对此问题的类似功能(我相信已经存在)。否则我会坚持现在的工作方式。

Answer 1

使用np.isin/np.in1d -

np.flatnonzero(np.isin(a,b))
# or np.flatnonzero(np.in1d(a,b))

对于已排序的 b -

，使用 np.searchsorted 可能性能更高

sidx = np.searchsorted(b,a)
sidx[sidx==len(b)] = len(b)-1
out = np.flatnonzero(b[sidx]==a)

如果 b 未排序，请对其进行排序，然后使用它代替 b。