numpy.diag_indices

numpy.diag_indices(n, ndim=2)

Функция diag_indices() возвращает индексы элементов главной диагонали квадратного массива заданного размера и размерности.

Словосочетание "квадратный массив" может быть обобщено на массивы большей размерности - не важно сколько осей у массива, главное что бы все они были одинаковой длины. А понятие диагональ массива означает что это на самом деле главная диагональ, т.е. диагональ начинается в точке отсчета индексов и заканчивается в точке их максимального значения.

Параметры:
n - целое положительное число.
Длина оси.
ndim - целое положительное число.
Количество осей, по умолчанию равно 2, что соответствует двумерным массивам; 3 - соответствует кубическим масивам; 4 - четырехмерным и т.д.
Возвращает:
tuple - кортеж массивов NumPy
кортеж массивов с индексами диагональных элементов.
Смотрите так же:
diag_indices_from, diag, diagflat, diagonal


Примеры

Извлечем диагональ массива:

>>> import numpy as np
>>> 
>>> a = np.arange(36).reshape(6, 6)
>>> 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, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35]])
>>> 
>>> d2 = np.diag_indices(6)
>>> d2
(array([0, 1, 2, 3, 4, 5]), array([0, 1, 2, 3, 4, 5]))
>>> 
>>> a[d2]
array([ 0,  7, 14, 21, 28, 35])

Пользуясь данным массивом индексов, мы можем присвоить диагональным элементам любое значение:

>>> a[d2] = -77
>>> a
array([[-77,   1,   2,   3,   4,   5],
       [  6, -77,   8,   9,  10,  11],
       [ 12,  13, -77,  15,  16,  17],
       [ 18,  19,  20, -77,  22,  23],
       [ 24,  25,  26,  27, -77,  29],
       [ 30,  31,  32,  33,  34, -77]])

Массив диагональных элементов может быть получен для массивов и большей размерности:

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

Что в свою очередь удобно для выполнения различных операций линейной алгебры, тензорного анализа или комбинаторики. Например, мы можем с легкостью сконструировать единичный трехмерный тензор:

>>> b = np.zeros((4, 4, 4))
>>> 
>>> b[d3] = 1
>>> b
array([[[1., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 1.]]])