numpy.diag_indices#

numpy.diag_indices(n, ndim=2)[source]#

返回访问数组主对角线的索引。

这返回一个索引元组，可用于访问形状为 (n, n, …, n) 且维度 a.ndim >= 2 的数组 a 的主对角线。对于 a.ndim = 2，这是通常的对角线，对于 a.ndim > 2，这是访问 a[i, i, ..., i] 的索引集，其中 i = [0..n-1]。

参数：
nint

返回的索引可用于的数组沿每个维度的尺寸。

ndimint，可选

维度数。

另请参阅

diag_indices_from

备注

版本 1.4.0 中的新功能。

示例

>>> import numpy as np

创建一组索引以访问 (4, 4) 数组的对角线

>>> di = np.diag_indices(4)
>>> di
(array([0, 1, 2, 3]), array([0, 1, 2, 3]))
>>> a = np.arange(16).reshape(4, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
>>> a[di] = 100
>>> a
array([[100,   1,   2,   3],
       [  4, 100,   6,   7],
       [  8,   9, 100,  11],
       [ 12,  13,  14, 100]])

现在，我们创建索引来操作一个 3D 数组

>>> d3 = np.diag_indices(2, 3)
>>> d3
(array([0, 1]), array([0, 1]), array([0, 1]))

并使用它将一个零数组的对角线设置为 1

>>> a = np.zeros((2, 2, 2), dtype=int)
>>> a[d3] = 1
>>> a
array([[[1, 0],
        [0, 0]],
       [[0, 0],
        [0, 1]]])
本页内容