# 创建一个以一维数组a为对角线，其余为0的二维数组
<<< a = np.arange(1, 4)
<<< a
<<< np.diag(a)
array([1, 2, 3])
array([[1, 0, 0],
       [0, 2, 0],
       [0, 0, 3]])

# 将二维数组b的对角线上的元素提取出来
b = np.arange(1, 10).reshape(3, 3)
b
np.diag(b)
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
array([1, 5, 9])