np.where()
Поэлементный выбор между x и y по условию (или поиск индексов).
Сигнатура
np.where(condition, x, y)С тремя аргументами возвращает массив, где истинным позициям соответствует x, ложным — y. С одним аргументом возвращает индексы истинных элементов.
import numpy as np
a = np.array([1, -2, 3, -4])
print(np.where(a > 0, a, 0)) # [1 0 3 0]
print(np.where(a > 0)) # (array([0, 2]),)