Поиск списка индексов из мастер-массива с использованием вторичного массива с неуникальными записями

Question

У меня есть мастер-массив длиной n номеров идентификаторов, которые применяются к другим аналогичным массивам с соответствующими данными для элементов в моем моделировании, которые принадлежат этим номерам идентификаторов (например, data[id]). Если бы я генерировал список номеров идентификаторов длиной m отдельно и нуждался в информации в массиве data для этих идентификаторов, каков наилучший способ получения списка индексов idx оригинала? массив идентификаторов для того, чтобы извлечь data[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

Есть ли лучший способ сделать это?

Answer 1

Текущий способ, которым вы делаете это, где поиск по всему массиву каждый раз.Вы можете сделать этот поиск O (1) вместо O (N), используя dict.Например, я использовал следующий метод:

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

и получил очень большое ускорение, которое будет еще лучше для больших массивов.Для размеров, которые вы имели в своем примере кода, я получил ускорение в 15 раз.Единственная проблема с моим кодом состоит в том, что если в a есть повторяющиеся элементы, то dict в настоящий момент будет указывать на последний экземпляр элемента, а с вашим методом - на первый экземпляр.Однако это можно исправить, если в фактическом использовании кода должны быть повторяющиеся элементы.

Answer 2

Я не уверен, есть ли способ сделать это автоматически в python, но вам, вероятно, лучше всего отсортировать два массива и затем сгенерировать вывод за один проход через b.Сложность этой операции должна составлять O(|a|*log|a|)+O(|b|*log|b|)+O(|b|) = O(|b|*log|b|) (при условии |b| > |a|).Я считаю, что ваша оригинальная попытка имеет сложность O(|a|*|b|), поэтому это должно обеспечить заметное улучшение для достаточно большого b.