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

利用Steven Rouk的解决方案，这里有一个方法来获得相等的子数组的索引：

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)

样品运行 -

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]))