Заменить элементы массива numpy на основе первого появления определенного значения

Question

Предположим, что имеется numpy 2D-массив следующего вида:

>>> x = np.array([[4,2,3,1,1], [1,0,3,2,1], [1,4,4,3,4]])
>>> x
array([[4, 2, 3, 1, 1],
       [1, 0, 3, 2, 1],
       [1, 4, 4, 3, 4]])

Моя цель - найти первое вхождение значения 4 в каждой строке и установить остальные элементы (включая этот элемент ) в этой строке на 0. Следовательно, после этой операции преобразованный массив должен выглядеть так:

>>> x_new
array([[0, 0, 0, 0, 0],
       [1, 0, 3, 2, 1],
       [1, 0, 0, 0, 0]])

Что такое pythoni c и оптимизированный способ сделать это? Я попытался использовать комбинацию np.argmax() и np.take(), но не смог достичь конечной цели.

Answer 1

Вы можете сделать это, используя кумулятивную сумму по столбцам (например, axis=1) и логическое индексирование:

n = 4
idx = np.cumsum(x == n, axis=1) > 0
x[idx] = 0

или, может быть, лучший способ - выполнить накопительное (логическое) или :

idx = np.logical_or.accumulate(x == n, axis=1)