Возвращать индексы общих элементов между двумя массивами

Question

У меня есть два массива, a1 и a2. Предположим, len(a2) >> len(a1), и что a1 является подмножеством a2.

Я бы хотел быстрый способ вернуть индексы a2 всех элементов в a1. Очевидно, что это требует много времени:

from operator import indexOf
indices = []
for i in a1:
    indices.append(indexOf(a2,i))

Конечно, это занимает много времени, когда а2 велико. Я мог бы также использовать numpy.where () вместо этого (хотя каждая запись в a1 будет появляться только один раз в a2), но я не уверен, что это будет быстрее. Я мог бы также пройти большой массив только один раз:

for i in xrange(len(a2)):
    if a2[i] in a1:
        indices.append(i)

Но я уверен, что есть более быстрый, более «тупой» способ - я просмотрел список «тупых» методов, но не могу найти ничего подходящего.

Большое спасибо заранее,

D

Answer 1

Как насчет

numpy.nonzero(numpy.in1d(a2, a1))[0]

Это должно быть быстро. Из моего базового тестирования это примерно в 7 раз быстрее, чем ваш второй фрагмент кода для len(a2) == 100, len(a1) == 10000, и только один общий элемент с индексом 45. Это предполагает, что и у a1, и a2 нет повторяющихся элементов.

Answer 2

как насчет:

wanted = set(a1)
indices =[idx for (idx, value) in enumerate(a2) if value in wanted]

Это должно быть O (len (a1) + len (a2)) вместо O (len (a1) * len (a2))

NB Я не знаю numpy, так что, возможно, есть более «numpythonic» способ сделать это, но так я бы сделал это на чистом питоне.

Answer 3

Очень похоже на @AlokSinghal, но вы получаете уже сплющенную версию.

numpy.flatnonzero(numpy.in1d(a2, a1))

Answer 4

index = in1d(a2,a1)
result = a2[index]

Answer 5

Пакет numpy_indexed (заявление об отказе: я его автор) содержит векторизованный эквивалент list.index; производительность должна быть аналогична принятому в настоящее время ответу, но в качестве бонуса она также дает вам явный контроль над отсутствующими значениями, используя «отсутствующий» kwarg.

import numpy_indexed as npi
indices = npi.indices(a2, a1, missing='raise')

Кроме того, он также будет работать с многомерными массивами, т. Е. Находить индексы одного набора строк в другом.