[FIXED] Numpy how to create mask "3D array contains 1D-value"

Issue

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.

Answered By - Warren Weckesser