最快重合矩阵

我有两个数组，我想计算一个巧合列表/数组。即所有索引 i、j 的列表，使得 a[i] == b[j]。这是我现在的代码：

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

np.array([[i, j] for i in range(a.size) for j in range(b.size) if a[i] == b[j]])

有没有更快的（也许是 numpy 驱动的）方法来做到这一点？

方法#1

一种方法是使用np.in1d -

m_a = np.in1d(a,b)
I = np.flatnonzero(m_a)
J = np.flatnonzero(np.in1d(b, a[m_a]))

样本输入、输出 -

In [367]: a
Out[367]: array([1, 2, 3, 4])

In [368]: b
Out[368]: array([3, 5, 6, 4])

In [370]: I
Out[370]: array([2, 3])

In [371]: J
Out[371]: array([0, 3])

方法#2

另一种直接但记忆繁重的方法是broadcasting -

I,J = np.nonzero(a[:,None] == b)

方法#3

对于输入数组中没有重复项的情况，我们可以使用np.searchsorted。这里有两种变体 - 一种用于排序a，另一个用于泛型a.

变体#1：对于已排序的a -

idx = np.searchsorted(a, b)
idx[idx==a.size] = 0
mask = a[idx] == b
I = np.searchsorted(a,b[mask])
J = np.flatnonzero(mask)

变体#2：对于这种通用变体情况，我们需要使用 argsort 索引a -

sidx = a.argsort()
a_sort = a[sidx]
idx = np.searchsorted(a_sort, b)
idx[idx==a.size] = 0
mask = a_sort[idx] == b
I = sidx[np.searchsorted(a_sort,b[mask])]
J = np.flatnonzero(mask)