一个numpy数组中元素的python位置与另一个数组中相等元素的位置

Question

我不仅需要值，还需要一个 numpy 数组中元素的位置，这些元素也出现在第二个 numpy 数组中，我也需要第二个数组中的位置。

这是我能做到的最好的例子:

>>> a=np.arange(0.,15.)
>>> a
array([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,
        11.,  12.,  13.,  14.])
>>> b=np.arange(4.,8.,.5)
>>> b
array([ 4. ,  4.5,  5. ,  5.5,  6. ,  6.5,  7. ,  7.5])
>>> [ (i,j) for (i,alem) in enumerate(a) for (j,blem) in enumerate(b) if alem==blem]
[(4, 0), (5, 2), (6, 4), (7, 6)]

有人有更快、特定于 numpy 或更“pythonic”的东西吗？

Answer 1

这是一个复杂度为 O((n+k)log(n+k)) 的解决方案(原始算法为 O(nk))np.unique

uniq, inv = np.unique(np.r_[a, b], return_inverse=True)
map = -np.ones((len(uniq),), dtype=int)
map[inv[:len(a)]] = np.arange(len(a))
bina = map[inv[len(a):]]
inds_in_b = np.where(bina != -1)[0]
elements, inds_in_a = b[inds_in_b], bina[inds_in_b]

或者您可以简单地将 a 排序为 O((n+k)log(k))

inds = np.argsort(a)
aso = a[inds]
bina = np.searchsorted(aso[:-1], b)
inds_in_b = np.where(b == aso[bina])[0]
elements, inds_in_a = b[inds_in_b], inds[bina[inds_in_b]]