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

Question

我有一个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

但这很难看，而且没有矢量化。对如何将其矢量化有什么建议吗？

谢谢

Answer 1

设置

x = np.array([[0. , 0. , 0.2, 0.2],
          [0.3, 0. , 0.3, 0. ],
          [0. , 0. , 0. , 0. ]])

使用logical_and与np.eye：

m = (x!=0).argmax(1)
x[~np.logical_and(x, np.eye(x.shape[1])[m])] = 0

输出：

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

使用这种方法会比其他两种建议的方法稍慢。

Answer 2

简单的

e =np.zeros(example.shape)
rows = np.arange(example.shape[0])
cols =  np.argmax(example != 0, 1)
e[rows, cols] = example[rows, cols]

Answer 3

这是一个矢量化的解决方案。诀窍是通过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. ]])