用数组的“元素号”替换数组中的元素（Python）

Question

我有一个数组，看起来像这样：

np.array([[0 , 5, 1], [0, 0, 3], [1, 7, 0]])

我想检查每个元素是否为非零，如果为非零，则用一个数字替换它，该数字跟踪它检查了多少个元素。也就是说，我希望最终产品看起来像

np.array([[0, 2, 3], [0, 0, 6], [7, 8, 0]])

其中第一行读取[0, 2, 3]，因为第二个元素是第二次检查的，通过了测试，然后被替换（依此类推）。有人能想出什么解决办法吗？我想numpy的索引在这里会非常有用。谢谢！你知道吗

Answer 1

如果需要修改初始数组-使用掩码数组的附加方法：

（来自IPython交互式控制台会话）

In [211]: arr = np.array([[0, 5, 1], [0, 0, 3], [1, 7, 0]])

In [212]: m = arr.nonzero()

In [213]: arr[m] = np.arange(1, arr.size+1).reshape(arr.shape)[m]

In [214]: arr
Out[214]: 
array([[0, 2, 3],
       [0, 0, 6],
       [7, 8, 0]])

Answer 2

你可以做：

np.where(a == 0, a, np.arange(a.size).reshape(a.shape) + 1)