Я хотел бы выполнить арифметику с k-й диагональю numpy.array. Мне нужны эти индексы. Например, что-то вроде:
>>> a = numpy.eye(2)
>>> a[numpy.diag_indices(a, k=-1)] = 5
>>> a
array([[ 1., 0.],
[ 5., 1.]])
К сожалению, diag_indices возвращает только индексы, составляющие главную диагональ, поэтому на данный момент я делаю:
a += numpy.diag([5], -1)
Но это не кажется таким красивым или надежным. :-)
Есть ли способ в numpy получить индексы, отличные от главной диагонали?
Итак, поскольку np.diag_indices() не имеет той же функциональности, что и np.triu_indices() для получения k-х диагоналей/треугольников, другой подход состоит в том, чтобы просто использовать np.eye(n,k) для построения матрицы nxn с единицами. на k-й диагонали, а затем используйте np.where, чтобы извлечь кортеж индексов, где расположены единицы.
Итак, мы можем сделать это, просто:
T = np.where(np.eye(5,k=-1) == 1)
Это дополнительное выделение матрицы глаза, которое в некоторых случаях может быть чрезмерным, но это легко сделать одной строкой.
Немного поздновато, но эта версия работает и для k = 0 (и не переделывает массивы, поэтому копировать не нужно).
def kth_diag_indices(a, k):
rows, cols = np.diag_indices_from(a)
if k < 0:
return rows[-k:], cols[:k]
elif k > 0:
return rows[:-k], cols[k:]
else:
return rows, cols
Вот способ:
Как это:
>>> import numpy as np
>>> rows, cols = np.indices((3,3))
>>> row_vals = np.diag(rows, k=-1)
>>> col_vals = np.diag(cols, k=-1)
>>> z = np.zeros((3,3))
>>> z[row_vals, col_vals]=1
>>> z
array([[ 0., 0., 0.],
[ 1., 0., 0.],
[ 0., 1., 0.]])
Индексы k-й диагонали a можно вычислить с помощью
def kth_diag_indices(a, k):
rowidx, colidx = np.diag_indices_from(a)
colidx = colidx.copy() # rowidx and colidx share the same buffer
if k > 0:
colidx += k
else:
rowidx -= k
k = np.abs(k)
return rowidx[:-k], colidx[:-k]
Демо:
>>> a
array([[ 0, 1, 2, 3, 4],
[ 5, 6, 7, 8, 9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])
>>> a[kth_diag_indices(a, 1)]
array([ 1, 7, 13, 19])
>>> a[kth_diag_indices(a, 2)]
array([ 2, 8, 14])
>>> a[kth_diag_indices(a, -1)]
array([ 5, 11, 17, 23])
Есть другое решение. Создайте матрицу E с помощью np.eye. Вы можете просто изменить главную диагональ, как показано ниже. Затем создайте еще две матрицы с параметром k в глазу. В конце просто сложите все матрицы вместе.
E = np.eye(5)
E = E*2
F = -1*np.eye(len(E),k=1)
G = -1*np.eye(len(E),k=-1)
E = E+F+G
print(E)
Где k задает расположение по диагонали от центра.
т.е. {k=0: центр по умолчанию, k=(-1): 1 ряд слева от центра, k=1: 1 ряд справа от центра}
Затем выполните арифметические действия, как обычно.
Ознакомьтесь с документацией здесь: np.diag().
In [3]: np.diag(np.arange(6), k=0)
Out[3]:
array([[0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0],
[0, 0, 2, 0, 0, 0],
[0, 0, 0, 3, 0, 0],
[0, 0, 0, 0, 4, 0],
[0, 0, 0, 0, 0, 5]])
In [4]: np.diag(np.arange(6), k=1)
Out[4]:
array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 1, 0, 0, 0, 0],
[0, 0, 0, 2, 0, 0, 0],
[0, 0, 0, 0, 3, 0, 0],
[0, 0, 0, 0, 0, 4, 0],
[0, 0, 0, 0, 0, 0, 5],
[0, 0, 0, 0, 0, 0, 0]])
In [5]: np.diag(np.arange(6), k=-1)
Out[5]:
array([[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 0],
[0, 0, 2, 0, 0, 0, 0],
[0, 0, 0, 3, 0, 0, 0],
[0, 0, 0, 0, 4, 0, 0],
[0, 0, 0, 0, 0, 5, 0]])
