Сопоставьте строки двух 2D-массивов и получите карту индексов строк, используя numpy

Question

Предположим, у вас есть два двумерных массива A и B, и вы хотите проверить, где находится строка A в B. Как вы делаете это наиболее эффективно, используя numpy?

Например,

a = np.array([[1,2,3],
              [4,5,6],
              [9,10,11]])

b = np.array([[4,5,6],
              [4,3,2],
              [1,2,3],
              [4,8,9]])
map = [[0,2], [1,0]]  # row 0 of a is at row index 2 of array B

Я знаю, как проверить, находится ли строка A в B, используя in1d ( тест на членство в двумерном массиве numpy ), но это не дает карту индексов.

Цель этой карты состоит в том, чтобы (наконец-то) объединить два массива на основе нескольких столбцов.
Конечно, можно делать этот ряд за строкой, но это становится очень неэффективным, поскольку мои массивы имеют форму (50 млн., 20).

Альтернативой может быть использование функции слияния pandas , но я бы хотел сделать это, используя только numpy.

Answer 1

Подход № 1

Вот один из них, основанный на views. Использует np.argwhere ( docs ) для возврата индексов элемента, которые удовлетворяют условию, в данном случае членству. -

def view1D(a, b): # a, b are arrays
    a = np.ascontiguousarray(a)
    b = np.ascontiguousarray(b)
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    return a.view(void_dt).ravel(),  b.view(void_dt).ravel()

def argwhere_nd(a,b):
    A,B = view1D(a,b)
    return np.argwhere(A[:,None] == B)

Подход № 2

Вот еще один, который будет O(n) и, следовательно, намного лучше по производительности, особенно на больших массивах -

def argwhere_nd_searchsorted(a,b):
    A,B = view1D(a,b)
    sidxB = B.argsort()
    mask = np.isin(A,B)
    cm = A[mask]
    idx0 = np.flatnonzero(mask)
    idx1 = sidxB[np.searchsorted(B,cm, sorter=sidxB)]
    return idx0, idx1 # idx0 : indices in A, idx1 : indices in B

Подход № 3

Другой O(n) один с использованием argsort() -

def argwhere_nd_argsort(a,b):
    A,B = view1D(a,b)
    c = np.r_[A,B]
    idx = np.argsort(c,kind='mergesort')
    cs = c[idx]
    m0 = cs[:-1] == cs[1:]
    return idx[:-1][m0],idx[1:][m0]-len(A)

Пример запускается с теми же входами, что и раньше -

In [650]: argwhere_nd_searchsorted(a,b)
Out[650]: (array([0, 1]), array([2, 0]))

In [651]: argwhere_nd_argsort(a,b)
Out[651]: (array([0, 1]), array([2, 0]))

Answer 2

Вы можете воспользоваться автоматической трансляцией:

np.argwhere(np.all(a.reshape(3,1,-1) == b,2))

, что приводит к

array([[0, 2],
       [1, 0]])

Примечание для чисел, которые вы можете заменить == на np.islclose()