如何在满足特定条件的情况下创建行索引的numpy矩阵

Question

如何将值的numpy矩阵转换为满足某个条件的行索引的numpy矩阵？

假设

A = array([[ 0.,  5.,  0.],[ 0.,  0.,  3.],[ 0.,  0.,  0.]])

如果我想在这里使用的条件 - 如果元素大于0，则替换为行索引 1，我该怎么办？因此，输出应该是

B = array([[0., 1., 0.],[0., 0., 2.],[0., 0., 0.]])

不确定我是否正确使用NP。谢谢。

Answer 1

使用 numpy.where

np.where(A>0, np.arange(1, A.shape[0]+1)[:, None], A)

array([[0., 1., 0.],
       [0., 0., 2.],
       [0., 0., 0.]])

---

或者您可以使用算术(如果您的值 比0少(：

(：

(A > 0) * np.arange(1, A.shape[0]+1)[:, None]