numpy.where()

numpy.where

numpy.where(condition[, x, y])

Return elements, either from x or y, depending on condition.

If only condition is given, return condition.nonzero().

Parameters:

condition : array_like, bool

When True, yield x, otherwise yield y.

x, y : array_like, optional

Values from which to choose. x and y need to have the same shape as condition.

Returns:

out