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

我有一个2-D 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. ]])