检查多维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和 中得到相同的 3 向量b。