Using 2d numpy mask np.where to address a 3d numpy array (pythonic??)

Question

Say I have an image array:

raster.shape => (3,100,100)

I generate a mask of all the places where red is saturated:

mask = np.where(raster[0,:,:] == 255)

I want to modify the slice of those locations... Which I can do like this:

raster[:,mask[0],mask[1]] = 0

Is there a more pythonic way... or is this the best bet?

Answer 1

I would use the mask directly, e.g.

raster[:, raster[0] == 255] = 0

which should be the equivalent of what you wrote.