numpy.diag()函数用法

numpy模块中有个函数diagonal，就是处理对角线的函数

numpy.diag(v,k=0)

以一维数组的形式返回方阵的对角线（或非对角线）元素，或将一维数组转换成方阵（非对角线元素为0）.两种功能角色转变取决于输入的v

v : array_like.
如果v是2D数组，返回k位置的对角线。

如果v是1D数组，返回一个v作为k位置对角线的2维数组。

k : int, optional
对角线的位置，大于零位于对角线上面，小于零则在下面。

当 v是1D数组时

import numpy as np
a=np.arange(4)
print('a= ')
print(a)
z=np.diag(a)
print('z= ')
print(z)

输出：

a=
[0 1 2 3]
z=
[[0 0 0 0]
[0 1 0 0]
[0 0 2 0]
[0 0 0 3]]

即以1D数组为对角线，其余元素为0的矩阵

当 v是2D数组时

import numpy as np
a=np.array([[1,2,3],[4,5,6],[7,8,9]])
print('a= ')
print(a)
z=np.diag(a)
print('z= ')
print(z)

输出：

a=
[[1 2 3]
[4 5 6]
[7 8 9]]
z=
[1 5 9]

输出对角线元素

设置K

import numpy as np
a=np.array([[1,2,3],[4,5,6],[7,8,9]])
print('a= ')
print(a)
z=np.diag(a)
print('z= ')
print(z)

z=np.diag(a,k=1)
print('k=1,z= ')
print(z)

z=np.diag(a,k=-1)
print('k=-1,z= ')
print(z)

输出：

a=
[[1 2 3]
[4 5 6]
[7 8 9]]

z=
[1 5 9]

k=1,z=
[2 6]

k=-1,z=
[4 8]

k=1，则输出对角线元素上移1的对角线
k=-1，则输出对角线元素下移1的对角线