Найти индексы общих значений в двух массивах

Я использую Python 2.7. У меня есть два массива: A и B. Чтобы найти индексы элементов в A, которые присутствуют в B, я могу сделать

A_inds = np.in1d(A,B) 

Я также хочу получить индексы элементов из B, которые присутствуют в A, т.е. индексы в B тех же перекрывающихся элементов, которые я нашел, используя приведенный выше код.

В настоящее время я снова запускаю одну строку:

 B_inds = np.in1d(B,A) 

но этот дополнительный расчет кажется, что он не должен быть лишним. Существует ли более эффективный вычислительный способ получения как A_inds и B_inds ?

Я открыт для использования методов списка или массива.

np.unique и np.searchsorted могут использоваться вместе для его решения –

 def unq_searchsorted(A,B): # Get unique elements of A and B and the indices based on the uniqueness unqA,idx1 = np.unique(A,return_inverse=True) unqB,idx2 = np.unique(B,return_inverse=True) # Create mask equivalent to np.in1d(A,B) and np.in1d(B,A) for unique elements mask1 = (np.searchsorted(unqB,unqA,'right') - np.searchsorted(unqB,unqA,'left'))==1 mask2 = (np.searchsorted(unqA,unqB,'right') - np.searchsorted(unqA,unqB,'left'))==1 # Map back to all non-unique indices to get equivalent of np.in1d(A,B), # np.in1d(B,A) results for non-unique elements return mask1[idx1],mask2[idx2] 

Проверка времени выполнения и проверка результатов –

 In [233]: def org_app(A,B): ...: return np.in1d(A,B), np.in1d(B,A) ...: In [234]: A = np.random.randint(0,10000,(10000)) ...: B = np.random.randint(0,10000,(10000)) ...: In [235]: np.allclose(org_app(A,B)[0],unq_searchsorted(A,B)[0]) Out[235]: True In [236]: np.allclose(org_app(A,B)[1],unq_searchsorted(A,B)[1]) Out[236]: True In [237]: %timeit org_app(A,B) 100 loops, best of 3: 7.69 ms per loop In [238]: %timeit unq_searchsorted(A,B) 100 loops, best of 3: 5.56 ms per loop 

Если два входных массива уже sorted и unique , повышение производительности будет значительным. Таким образом, функция решения упростит –

 def unq_searchsorted_v1(A,B): out1 = (np.searchsorted(B,A,'right') - np.searchsorted(B,A,'left'))==1 out2 = (np.searchsorted(A,B,'right') - np.searchsorted(A,B,'left'))==1 return out1,out2 

Последующие тесты времени выполнения –

 In [275]: A = np.random.randint(0,100000,(20000)) ...: B = np.random.randint(0,100000,(20000)) ...: A = np.unique(A) ...: B = np.unique(B) ...: In [276]: np.allclose(org_app(A,B)[0],unq_searchsorted_v1(A,B)[0]) Out[276]: True In [277]: np.allclose(org_app(A,B)[1],unq_searchsorted_v1(A,B)[1]) Out[277]: True In [278]: %timeit org_app(A,B) 100 loops, best of 3: 8.83 ms per loop In [279]: %timeit unq_searchsorted_v1(A,B) 100 loops, best of 3: 4.94 ms per loop 

Простая многопроцессорная реализация даст вам немного больше скорости:

 import time import numpy as np from multiprocessing import Process, Queue a = np.random.randint(0, 20, 1000000) b = np.random.randint(0, 20, 1000000) def original(a, b, q): q.put( np.in1d(a, b) ) if __name__ == '__main__': t0 = time.time() q = Queue() q2 = Queue() p = Process(target=original, args=(a, b, q,)) p2 = Process(target=original, args=(b, a, q2)) p.start() p2.start() res = q.get() res2 = q2.get() print time.time() - t0 >>> 0.21398806572 

unq_searchsorted(A,B) Divakar занял 0.271834135056 секунд на моей машине.

Interesting Posts