NumPy第k个对角线指数

Question

我想做一个numpy.array的第k个对角线的算术.我需要那些指数.例如,类似于:

>>> 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有没有办法获得除主要对角线以外的指数？

Answer 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

Answer 2

这是一种方式:

1. 创建索引值数组.
2. 获取所需的daigonal索引值.
3. 而已!:)

像这样:

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