numpy матричная диагональная заливка с переменным значением

Question

Я ищу эффективный способ записать переменное значение по диагонали матрицы. например ..

array([[1, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 1, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 1, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 1, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 0, 0, 1]])

Answer 1

Попробуйте np.eye с некоторыми нарезками:

import numpy as np
a = np.eye(9)
a[1::2] = 0
print(a)

Вывод:

array([[1., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 1., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 1., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0., 0., 0., 1.]])

Answer 2

np.diag использует трюк присваивания flat. Для этого случая его можно упростить с помощью:

a=np.zeros((9,9))
a.flat[10::20]=1  

Это оказывается значительно быстрее, чем предложенные альтернативы (пока):

In [108]: timeit a=np.zeros((9,9)); a.flat[10::20]=1                                             
1.5 µs ± 10.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [110]: timeit np.diag(([1,0]*((k+1)//2))[:k])                                                 
10.1 µs ± 29.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

In [111]: timeit a=np.eye(9); a[1::2]=0                                                          
6.74 µs ± 20 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Подход eye устанавливает целое строк до 0, но поскольку это делается с помощью широковещания, это относительно быстро.

Answer 3

Это например:

k = 9
np.diag(([1,0]*((k+1)//2))[:k])

или

np.diag([1,0]*int(np.ceil(k/2)))