numpy.diag_indices
numpy.diag_indices(n, ndim=2)
Функция diag_indices() возвращает индексы элементов главной диагонали квадратного массива заданного размера и размерности.
Словосочетание "квадратный массив" может быть обобщено на массивы большей размерности - не важно сколько осей у массива, главное что бы все они были одинаковой длины. А понятие диагональ массива означает что это на самом деле главная диагональ, т.е. диагональ начинается в точке отсчета индексов и заканчивается в точке их максимального значения.
-
- n - целое положительное число.
- Длина оси.
- ndim - целое положительное число.
- Количество осей, по умолчанию равно 2, что соответствует двумерным массивам; 3 - соответствует кубическим масивам; 4 - четырехмерным и т.д.
-
- tuple - кортеж массивов NumPy
- кортеж массивов с индексами диагональных элементов.
Примеры
Извлечем диагональ массива:
>>> 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.]]])