найти показатели ndarray по сравнению с ndarray

Question

У меня есть два несортированных ndarrays со следующей структурой:

a1 = np.array([0,4,2,3],[0,2,5,6],[2,3,7,4],[6,0,9,8],[9,0,6,7])
a2 = np.array([3,4,2],[0,6,9])

Я хотел бы найти все индексы a1, где каждая строка a2 находится в a1, а также внутри позиции a1:

result = [[0,[3,1,2]],[2,[1,3,0]],[3,[1,0,2]],[4,[1,2,0]]

В этом примере a2 [0] находится в a1 в положении 0 и 2 в положении a1 в 3,1,2 и 1,3,0.Для a2 [1] в позициях 3 и 4 в пределах a1 в позициях 1,0,2 и 1,2,0.

Каждая строка a2 появляется дважды в a1.а1 имеет минимум 1 миллионряды, а2 около 10000.Таким образом, алгоритм должен быть также довольно быстрым (если возможно).

До сих пор я думал об этом подходе:

big_res = []
for r in xrange(len(a2)):
    big_indices = np.argwhere(a1 == a2[r])
    small_res = []
    for k in xrange(2):
        small_indices = [i for i in a2[r] if i in a1[big_indices[k]]]
        np.append(small_res, small_indices)
    combined_res = [[big_indices[0],small_res[0]],[big_indices[1],small_res[1]]]
    np.append(big_res, combined_res)

Answer 1

Использование numpy_indexed , (заявление об отказе: я его автор), то, что я считаю трудной частью, может быть эффективно записано следующим образом:

import numpy_indexed as npi

a1s = np.sort(a1, axis=1)
a2s = np.sort(a2, axis=1)
matches = np.array([npi.indices(a2s, np.delete(a1s, i, axis=1), missing=-1) for i in range(4)])
rows, cols = np.argwhere(matches != -1).T
a1idx = cols
a2idx = matches[rows, cols]
# results.shape = [len(a2), 2]
result = npi.group_by(a2idx).split_array_as_array(a1idx)

Это дает только совпаденияэффективно;не относительные порядки.Но как только у вас есть совпадения, вычисление относительных порядков должно быть простым в линейном времени.

РЕДАКТИРОВАТЬ: и некоторый код сомнительной плотности, чтобы получить ваши относительные порядки:

order = npi.indices(
    (np.indices(a1.shape)[0].flatten(), a1.flatten()),
    (np.repeat(result.flatten(), 3),    np.repeat(a2, 2, axis=0).flatten())
).reshape(-1, 2, 3) - result[..., None] * 4