根据特定值的第一次出现替换 numpy 数组的元素

Question

假设有一个 numpy 二维数组如下：

>>> 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]])

这样做的pythonic和优化方式是什么？我尝试了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)