numpy.ma.where

numpy.ma.where(condition | x, y)

Returns a (subclass of) masked array, shaped like condition, where the elements are x when condition is True, and y otherwise. If neither x nor y are given, returns a tuple of indices where condition is True (a la condition.nonzero()).

Parameters:

condition : {var}

The condition to meet. Must be convertible to an integer array.

x : {var}, optional

Values of the output when the condition is met

y : {var}, optional

Values of the output when the condition is not met.