numpy diag

NumPy的`diag`函数用于提取或构造一个矩阵的对角线元素。它有两种用法：

1. 提取对角线元素：可以传入一个矩阵作为参数，返回该矩阵的一维数组，其中包含原矩阵的主对角线元素。例如：

import numpy as np

a = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]])

diag_a = np.diag(a)
print(diag_a)  # 输出 [1 5 9]

2. 构造对角矩阵：可以传入一个一维数组作为参数，返回一个以该数组为主对角线元素的对角矩阵。例如：

import numpy as np

arr = np.array([1, 2, 3])
diag_arr = np.diag(arr)
print(diag_arr)

输出为：

[[1 0 0]
 [0 2 0]
 [0 0 3]]

希望这个解答能够帮到您！如果您还有其他问题，请随时提问。