Python сравнение два 3 мерных Numpy массивов

Question

У меня есть два Numpy массива:

A.shape = (nA,x,y) 

и

B.shape = (nB,x,y). 

Я хочу, чтобы найти все подмассив такие, что

A(i,:,:) == B(j,:,:). 

Я знаю, что может напишите двойной цикл и используйте

np.array_equal(A(i,:,:),B(j,:,:) 

Однако существует ли более эффективный метод?

Answer 1

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

subarrays = [x for x in A if x in B] 

Если вы хотите только индексы, а не хранить весь подмассив, вы можете сделать:

indices = [x[0] for x in enumerate(A) if x[1] in B] 

Answer 2

Используя решение Стивена Rõuk, вот способ получить индексы для подрешетки, которые равны :

indicesForMatches = [(i,j) for i,subArrayOfA in enumerate(A) for j,subArrayOfB in enumerate(B) if np.array_equal(subArrayOfA,subArrayOfB)] 

Answer 3

Вы можете использовать NumPy broadcasting для векторизованного решения, например, так -

mask = ((A[:,None,:,:] == B).all(2)).all(2) 
A_idx,B_idx = np.where(mask) 

Вы можете использовать reshaping, чтобы избежать двойных .all() использований и получить маску, как так -

mask = (A.reshape(A.shape[0],1,-1) == B.reshape(B.shape[0],-1)).all(-1) 

Sample пробег -

In [41]: # Setup input arrays and force some indices to be same between A and B 
    ...: nA = 4 
    ...: nB = 5 
    ...: x = 3 
    ...: y = 2 
    ...: 
    ...: A = np.random.randint(0,9,(nA,x,y)) 
    ...: B = np.random.randint(0,9,(nB,x,y)) 
    ...: 
    ...: A[2,:,:] = B[1,:,:] 
    ...: A[3,:,:] = B[4,:,:] 
    ...: 

In [42]: mask = ((A[:,None,:,:] == B).all(2)).all(2) 
    ...: A_idx,B_idx = np.where(mask) 
    ...: 

In [43]: A_idx, B_idx 
Out[43]: (array([2, 3]), array([1, 4])) 

In [44]: mask = (A.reshape(A.shape[0],1,-1) == B.reshape(B.shape[0],-1)).all(-1) 
    ...: A_idx,B_idx = np.where(mask) 
    ...: 

In [45]: A_idx, B_idx 
Out[45]: (array([2, 3]), array([1, 4]))