NumPy k-ый диагональный индекс

Я бы хотел сделать арифметику с 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, чтобы получить индексы, отличные от основной диагонали?

Ответ 1

Немного поздно, но эта версия также работает для k = 0 (и не изменяет массивы, поэтому не нужно делать копию). Это работает для массивов, которые также не квадратные.

def kth_diag_indices(a, k):
     rows = np.arange(a.shape[0])
     cols = np.arange(a.shape[1])
     if k >= 0:
         return rows[:a.shape[1]-k], cols[k:a.shape[0]+k]
     else:
         return rows[-k:a.shape[1]], cols[:a.shape[0]+k]

Ответ 2

Вот путь:

Создание массивов значений индекса.
Получите нужные значения параметров дайгана.
Вот оно!:)

Вот так:

>>> 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.]])

Ответ 3

Индексы 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])

Ответ 4

Используйте `numpy.diag(v, k=0)`

Где 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]])