NumPy 第 k 个对角线索引

Posted 2023-02-14

【中文标题】NumPy 第 k 个对角线索引

【英文标题】：NumPy k-th diagonal indices

【发布时间】：2012-06-11 03:28:50

【问题描述】：

我想用 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 中有没有办法获取除主对角线以外的索引？

【参考方案1】：

所以由于 np.diag_indices() 没有与 np.triu_indices() 相同的功能来获取第 k 个对角线/三角形，另一种方法是只使用 np.eye(n,k) 来构造一个 nxn在第 k 个对角线上具有 1 的矩阵，然后使用 np.where 提取 1 所在位置的索引元组。

所以我们可以这样做：

T = np.where(np.eye(5,k=-1) == 1)

这是对眼睛矩阵的额外分配，在某些情况下可能会过度，但它是一种简单的单线。

【讨论】：

虽然这在理论上可以回答这个问题，但如果您添加一些文字来解释为什么您的回答在提问者 8 年前已经接受了一个答案之后对每个读者来说应该是一个不错的选择，那就更好了。

【参考方案2】：

还有另一种解决方案。使用 np.eye 创建一个矩阵 E。您可以像下面这样更改主对角线。然后，用 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)

【参考方案3】：

有点晚了，但是这个版本也适用于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

【参考方案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]])

【讨论】：

是的，我知道如何建立一个新的对角矩阵。但是当我需要修改 一个现有的 矩阵时，您上面的方法并不完全适用。

为什么不呢？在第 k 个对角线上创建一个矩阵，并使用它来索引您的原始矩阵。

@Stefan 这不适用于非方形数组。请使用其他答案。

【参考方案5】：

a 的第 k 条对角线的索引可以用

计算

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

【参考方案6】：

这是一种方法：

创建索引值数组。 获取所需的对角索引值。 就是这样！ :)

像这样：

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