檢查多維numpy數組中兩個軸的相等性

Question

我得到了一個 3 維形狀 (n,m,k) numpy 數組。 我想將其視為包含向量的二維矩陣，即具有大小為 k 的向量的 nxm 矩陣。 我現在想檢查兩個這樣的形狀數組 (n,m,k) 是否第一個數組中的條目 (x,y,:) 等於第二個數組中的 (x,y,:) 。 有沒有一種方法可以在不使用循環的情況下在 numpy 中做到這一點？

我想過類似 A == B 以第一和第二軸為條件的東西。

Answer 1

您可以將條件和ndarray.all與axis一起使用：

a = np.arange(27).reshape(3,3,3)
b = np.zeros_like(a)
b[0,1,2] = a[0,1,2]
b[1,2,0] = a[1,2,0]
b[2,1,:] = a[2,1,:] # set to the same 3-vector at n=2, m=1

(a == b).all(axis=2) # check whether all elements of last axis are equal
# array([[False, False, False],
#        [False, False, False],
#        [False,  True, False]])

如您所見，對於n=2和m=1我們在a和b得到相同的 3-vector。