Замена элементов в массиве его «номером элемента» (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

В случае, если вам нужно изменить исходный массив - дополнительный подход с использованием mask array:

(из сеанса интерактивной консоли 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]])