Найти индексы совпадающих строк в двух двухмерных массивах

Предположим, что у меня два двухмерных массива:

array([[3, 3, 1, 0],
       [2, 3, 1, 3],
       [0, 2, 3, 1],
       [1, 0, 2, 3],
       [3, 1, 0, 2]], dtype=int8)

array([[0, 3, 3, 1],
       [0, 2, 3, 1],
       [1, 0, 2, 3],
       [3, 1, 0, 2],
       [3, 3, 1, 0]], dtype=int8)

Некоторые строки в каждом массиве имеют соответствующую строку, которая соответствует по значению (но не обязательно по индексу) в другом массиве, а некоторые нет.

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

(0,4)
(2,1)
(3,2)
(4,3)

Ответ 1

Это все numpy решение - не обязательно обязательно лучше, чем итеративный Python. Он все равно должен смотреть на все комбинации.

In [53]: np.array(np.all((x[:,None,:]==y[None,:,:]),axis=-1).nonzero()).T.tolist()
Out[53]: [[0, 4], [2, 1], [3, 2], [4, 3]]

Промежуточный массив (5,5,4). np.all сводит его к:

array([[False, False, False, False,  True],
       [False, False, False, False, False],
       [False,  True, False, False, False],
       [False, False,  True, False, False],
       [False, False, False,  True, False]], dtype=bool)

Остальное - это просто извлечение индексов, где это True

В сырых тестах, это время на 47,8 мы; другой ответ с помощью словаря L1 на 38.3 нам; и третий с двойной петлей на 496 нас.

Ответ 2

Я не могу придумать многозначный способ сделать это, но вот что я буду делать с регулярными списками:

>>> L1= [[3, 3, 1, 0],
...        [2, 3, 1, 3],
...        [0, 2, 3, 1],
...        [1, 0, 2, 3],
...        [3, 1, 0, 2]]
>>> L2 = [[0, 3, 3, 1],
...        [0, 2, 3, 1],
...        [1, 0, 2, 3],
...        [3, 1, 0, 2],
...        [3, 3, 1, 0]]
>>> L1 = {tuple(row):i for i,row in enumerate(L1)}
>>> answer = []
>>> for i,row in enumerate(L2):
...   if tuple(row) in L1:
...     answer.append((L1[tuple(row)], i))
... 
>>> answer
[(2, 1), (3, 2), (4, 3), (0, 4)]

Ответ 3

Вы можете использовать трюк типа данных void для использования 1D-функций в строках ваших двух массивов. a_view и b_view являются 1D векторами, каждая запись представляет собой полную строку. Затем я решил отсортировать массив и использовать np.searchsorted, чтобы найти элементы другого массива в этом. Если массив, который мы сортируем, имеет длину m, а второй имеет длину n, сортировка занимает время m * log(m), а двоичный поиск, который np.searchsorted, занимает время n * log(m), всего (n + m) * log(m). Поэтому вы хотите отсортировать самый короткий из двух массивов:

def find_rows(a, b):
    dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))

    a_view = np.ascontiguousarray(a).view(dt).ravel()
    b_view = np.ascontiguousarray(b).view(dt).ravel()

    sort_b = np.argsort(b_view)
    where_in_b = np.searchsorted(b_view, a_view,
                                 sorter=sort_b)
    where_in_b = np.take(sort_b, where_in_b)
    which_in_a = np.take(b_view, where_in_b) == a_view
    where_in_b = where_in_b[which_in_a]
    which_in_a = np.nonzero(which_in_a)[0]
    return np.column_stack((which_in_a, where_in_b))

С a и b ваши два массива образцов:

In [14]: find_rows(a, b)
Out[14]: 
array([[0, 4],
       [2, 1],
       [3, 2],
       [4, 3]], dtype=int64)

In [15]: %timeit find_rows(a, b)
10000 loops, best of 3: 29.7 us per loop

В моей системе словарь приближается быстрее примерно на 22 us для ваших тестовых данных, но с массивами 1000x4 этот метод numpy примерно в 6 раз быстрее, чем чистый Python (483 us vs 2.54 ms).