Check equality of two axes in multidiimensional numpy array

Question

I am given a 3-dimensional shape(n,m,k) numpy array. I'd like to view this as a 2-dimensional matrix containing vectors, i.e. a nxm matrix with a vector of size k. I'd now like to check for two such arrays of shape (n,m,k) wheter entry (x,y,:) in the first array is equal to (x,y,:) in the second array. Is there a method to do this in numpy without using loops?

I'd thought about something like A == B conditioned on the first and second axis.

Answer 1

You can use a condition, and ndarray.all together with 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]])

As you can see, for n=2 and m=1 we get the same 3-vector in a and b.