Как создать матрицу Numpy индекса строки, где выполняется определенное условие?

Question

Как преобразовать цифровую матрицу значений в цифровую матрицу индексов строк, где выполняется определенное условие?

Скажем,

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.where.Спасибо.

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]