I have two arrays: 3D-array ndarray(N,M,3) and 1D-array ndarray(3). How to create mask N×M by comparing last axis of 3D array with 1D array?
For now, I can do this with:
mask = (A[:,:,0] == B[0]) & (A[:,:,1] == B[1]) & (A[:,:,2] == B[2])
But there must be a shorter solution.
>Solution :
mask = (A == B).all(axis=2)
For more information about how this works, search numpy broadcasting.