Python numpy diag_indices用法及代碼示例

本文簡要介紹 python 語言中 numpy.diag_indices 的用法。

用法:

numpy.diag_indices(n, ndim=2)

返回索引以訪問數組的主對角線。

這將返回一個索引元組，可用於訪問數組的主對角線a和a.ndim >= 2尺寸和形狀(n，n，...，n)。為了a.ndim = 2這是通常的對角線，對於a.ndim > 2這是要訪問的索引集a[i, i, ..., i]為了i = [0..n-1].

參數：

n： int

可以使用返回的索引的數組的每個維度的大小。

ndim： 整數，可選

維數。

注意：

例子：

創建一組索引以訪問 (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]])

現在，我們創建索引來操作 3-D 數組：

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