python之保留第一个非零元素，将所有其他元素设置为 0

我有一个二维 NumPy 数组，如下所示:

array([[0. , 0. , 0.2, 0.2], 
       [0.3, 0. , 0.3, 0. ]]) 

我想修改它，使每一行都由全 0 组成，除了第一个非零条目。如果一开始都是 0，我们什么都不做。

我可以这样做:

example = np.array([[0,0, 0.2, 0.2], [0.3, 0, 0.3, 0]]) 
my_copy = np.zeros_like(example) 
for i, row in enumerate(example): 
    for j, elem in enumerate(row): 
        if elem > 0: 
            my_copy[i, j] = elem 
            break 

但这很丑陋而且没有向量化。关于如何对其进行矢量化有什么建议吗？

谢谢!

请您参考如下方法:

这是一个矢量化的解决方案。诀窍是通过 bool 转换和 argmax 计算您的第一个非零条目。

import numpy as np 
 
A = np.array([[0. , 0. , 0.2, 0.2], 
              [0.3, 0. , 0.3, 0. ], 
              [0. , 0. , 0. , 0. ]]) 
 
res = np.zeros(A.shape) 
idx =  np.arange(res.shape[0]) 
args = A.astype(bool).argmax(1) 
res[idx, args] = A[idx, args] 
 
print(res) 
 
array([[ 0. ,  0. ,  0.2,  0. ], 
       [ 0.3,  0. ,  0. ,  0. ], 
       [ 0. ,  0. ,  0. ,  0. ]])