Skip to main content
 首页 » 编程设计

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

2025年05月04日12linjiqin

我不仅需要值,还需要一个 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”的东西吗?

请您参考如下方法:

这是一个复杂度为 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]]