NumPy中的np.diag函数介绍

np.diag()是NumPy中的一个函数，它用于将一维数组转换为方阵或从一个方阵中提取对角线:

将一维数组转换为方阵

import numpy as np

# 将一维数组转换为方阵
a = [1, 2, 3]
b = np.diag(a)
print(b)

输出结果为：

array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]])

此时，生成的方阵为一个3行3列的方阵，且对角线上的元素为[1,2,3]。

提取方阵中的对角线元素

import numpy as np

# 提取方阵中的对角线元素
a = np.array([[1, 0, 0], [0, 2, 0], [0, 0, 3]])
b = np.diag(a)
print(b)

输出结果为：

array([1, 2, 3])

此时，提取出的元素为方阵的对角线元素。

np.diag函数的参数

np.diag()函数的参数有两种，一种是一维数组，用于生成方阵；另一种是方阵，用于提取对角线元素。下面的例子中我们使用一个嵌套列表表示的方阵：

import numpy as np

# 将一维数组转换为方阵
a = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
b = np.diag(a)
print(b)

# 提取方阵中的对角线元素
a = np.array([[1, 0, 0], [0, 2, 0], [0, 0, 3]])
b = np.diag(a)
print(b)

输出结果与之前相同。

矩阵运算

np.diag()函数还可以进行矩阵运算，例如：

import numpy as np

# 矩阵运算
a = np.diag([1,2,3])
b = np.array([[1,2,3], [4,5,6], [7,8,9]])
print(np.dot(a, b))

此时将对角线元素[1,2,3]转换为3行3列的方阵，与3行3列的b矩阵进行矩阵乘法运算，得到一个新的3行3列的矩阵。