我有一个二维 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. ]])
关于python - 保留第一个非零元素,将所有其他元素设置为 0,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/51731815/