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