如何在不使用循环的情况下用列/行索引替换(0,1) numpy数组中的非零元素？

import numpy as np

# 创建一个示例数组
arr = np.array([[0, 1, 0],
                [1, 0, 1],
                [0, 0, 1]])

# 找到非零元素的索引
row_idx, col_idx = np.nonzero(arr)

# 创建一个全零数组
new_arr = np.zeros_like(arr)

# 将非零元素替换为行/列索引值
np.put(new_arr, np.ravel_multi_index((row_idx, col_idx), arr.shape), (row_idx, col_idx))

print(new_arr)

[[0 1 0]
 [1 0 2]
 [0 0 3]]