最快的重合矩阵 [英] Fastest coincidence matrix

问题描述

我有两个数组，我想计算一个巧合的列表/数组.即，列出所有索引i，j，以便a [i] == b [j].现在这是我的代码:

I have two arrays and I want to compute a list/array of coincidences. That is, a list of all the indices i, j so that a[i] == b[j]. This is my code now:

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驱动)来做到这一点?

Is there a faster (maybe numpy-powered) way to do this?

推荐答案

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

Approach #3

For the case where we have no duplicates within the input arrays, we could use np.searchsorted. There are two variants here - One for sorted a, and another for a generic 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-

Variant #2 : For this generic variant case, we need to use argsort indices of 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)

这篇关于最快的重合矩阵的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！