用“元素編號”替換數組中的元素（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

你可以做：

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

Answer 2

如果您需要修改初始數組 - 使用掩碼數組的附加方法：

（來自 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]])