numpy_diag函数

  diag的函数原型如下：

numpy.diag(v, k=0)

• 如果v是一维数组，则函数返回以v为对角线的对角矩阵。
• 如果v是二维数组，则函数返回v的对角线。

import numpy as np

a = np.arange(9).reshape((3, 3))
print("a =", a)
b = np.diag(a)
print("b =", b)
c = np.diag(b)
print("c =", c)

执行结果：

a = [[0 1 2]
     [3 4 5]
     [6 7 8]]
b = [0 4 8]
c = [[0 0 0]
     [0 4 0]
     [0 0 8]]