python - 最快重合矩阵

Question

我有两个数组，我想计算一个巧合列表/数组。也就是说，所有索引 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 支持的)方法来做到这一点？

Answer 1

方法#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

另一种直接但占用大量内存的方法是使用广播 -

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:对于这种通用变体情况，我们需要使用 a 的 argsort 索引 -

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)