Python 使用具有非唯一项的辅助数组从主数组查找索引列表

我有一个id号长度为n的主数组，该数组应用于其他类似数组，并具有我的模拟中属于这些id号的元素的相应数据（例如，

data[id]

）。如果我要单独生成长度为m的id编号列表，并且需要这些id的

数据

数组中的信息，那么为了提取

数据[idx]

，获取原始id数组的索引列表

idx

的最佳方法是什么？也就是说，考虑到：

a=numpy.array([1,3,4,5,6])      # master array
b=numpy.array([3,4,3,6,4,1,5])  # secondary array

我想生成

idx=numpy.array([1,2,1,4,2,0,3])

数组

a

通常按顺序排列，但不是必需的。另外，数组

b

肯定会有重复，并且不会以任何顺序排列

我目前的做法是：

idx=numpy.array([numpy.where(a==bi)[0][0] for bi in b])

我使用以下测试对其进行计时：

a=(numpy.random.uniform(100,size=100)).astype('int')
b=numpy.repeat(a,100)
timeit method1(a,b)

10 loops, best of 3: 53.1 ms per loop

---

有更好的方法吗？

我不确定python中是否有自动执行此操作的方法，但您最好对两个数组进行排序，然后通过一次

b

生成输出。该操作的复杂性应该是

O（|a |*log | a |）+O（|b |*log | b |）+O（| b |）=O（| b |*log | b |）

（假设

| b |>a

）。我相信您最初的尝试具有复杂性

O（|a |*|b |）

，因此这将为足够大的

b

提供一个明显的改进，您当前的方法是每次搜索整个a数组。您可以使用dict查找O（1）而不是O（N）。例如，我使用了以下方法：

def method2(a,b):
    tmpdict = dict(zip(a,range(len(a))))
    idx = numpy.array([tmpdict[bi] for bi in b])

---

并且获得了非常大的加速，这对于更大的阵列来说会更好。对于示例代码中的大小，我得到了15倍的加速。我的代码唯一的问题是，如果

a

中有重复的元素，那么dict当前将指向元素的最后一个实例，而使用您的方法，它将指向第一个实例。但是，如果在代码的实际使用中有重复的元素，则可以纠正这种情况

非常好。非常感谢。这很有效，因为

a

中没有重复的元素，因为它是唯一id号的列表。@fideli，这是我猜的，但是你的随机数示例不排除重复。很高兴我能帮忙。