如何将 Python NumPy 数组中的值替换为在另一个数组中找到的那些值的索引？

编程问答 2022-06-01

问题描述

我有一个 n*m 数组“a”，另一个一维数组“b”，如下所示：

a = array([[ 51,30,20,10],[ 10,32,65,77],[ 15,77,30]])

b = array([10,15,51,77])

我想将“a”中的所有元素替换为该元素所在的“b”的相应索引。在上述情况下，我希望输出为：

a = array([[ 5,3,2,0],[ 0,4,6,7],[ 1,7,3]])

请注意，在实际应用中，我的数组很大，超过 30k 个元素和数千个元素。我试过 for 循环，但这些需要很长时间来计算。我也尝试过类似的迭代方法，并使用 list.index() 来获取索引，但这也需要太多时间。

谁能帮我首先确定出现在“b”中的“a”元素的“b”索引，然后构造更新的“a”数组？

谢谢。

解决方法

如果 a,b 的最小/最大元素形成一个小范围（或至少小到足以放入 RAM），则可以使用查找表非常快速地完成此操作：

a = np.array([[51,30,20,10],[10,32,65,77],[15,77,30]])
b = np.array([10,15,51,77])

lo = min(a.min(),b.min())
hi = max(a.max(),b.max())
lut = np.zeros(hi - lo + 1,dtype=np.int64)
lut[b - lo] = np.arange(len(b))

那么：

>>> a_indices = lut[a - lo]
>>> a_indices
array([[5,3,2,0],[0,4,6,7],[1,7,3]])
,

这只是作为答案发布的，因为评论太长了。它支持上面发布的 orlp 的解决方案。 Numpy 的 vectorize 避免了显式循环，但这显然不是最好的方法。 请注意，Numpy 的 searchsorted 只能在 b 已排序时如图所示应用。

import timeit
import numpy as np

a = np.random.randint(1,100,(1000,1000))
b = np.arange(0,1000,1)

def o1():
    lo = min(a.min(),b.min())
    hi = max(a.max(),b.max())
    lut = np.zeros(hi - lo + 1,dtype=np.int64)
    lut[b - lo] = np.arange(len(b))
    a2 = lut[a - lo]
    return a2 

def o2():
    a2 = a.copy()
    fu = np.vectorize(lambda i: np.place(a2,a2==b[i],i))
    fu(np.arange(0,len(b),1))

print(timeit.timeit("np.searchsorted(b,a)",globals=globals(),number=2))
print(timeit.timeit("o1()",number=2))
print(timeit.timeit("o2()",number=2))

印刷品

0.061956800000189105
0.012765400000716909
2.220097600000372
arraysarraysarraysindexingindexingindiceslargenumberpython