np.diag 提取主对角线最简单,但输入为一维数组时会生成对角矩阵而非提取对角线;需确认 arr.ndim == 2 或先用 np.atleast_2d(arr);k 参数控制偏移,有效范围为 -(m-1) ≤ k ≤ n-1。

直接用 np.diag 提取主对角线最简单,但别默认它只返回一维数组
调用 np.diag 时,如果输入是二维数组(如 ndarray),它默认提取主对角线并返回一维 ndarray;但如果输入是一维数组,它反而会生成对角矩阵。这个行为容易混淆,尤其在链式操作中意外把向量当矩阵传进去,结果得到一个方阵而非预期的对角元素列表。
- 确认输入维度:
arr.ndim == 2才能安全提取对角线 - 若不确定,先用
np.atleast_2d(arr)包裹再调用np.diag - 注意返回值类型:总是
ndarray,不是 Python list,切片或索引仍需用 NumPy 方式
np.diag 的 k 参数控制偏移,正数向上、负数向下
k 是关键偏移参数:当 k=0(默认)取主对角线;k=1 取主对角线上方第一条(即第 0 行第 1 列开始);k=-2 取主对角线下方第二条(即第 2 行第 0 列开始)。实际提取范围受矩阵形状限制,超出部分自动截断,不报错也不补零。
- 对于形状为
(m, n)的矩阵,有效k范围是-(m-1) -
np.diag(arr, k=5)在 4×4 矩阵上返回空数组array([]),不是错误 - 若需固定长度结果(如补零),得手动处理,
np.diag不提供填充选项
想同时取多条对角线?np.diag 不行,得换思路
np.diag 每次只能取一条对角线。要批量提取(比如主对角线 + 上下各一条),不能靠多次调用再拼接——效率低且难对齐。更稳的方式是用高级索引构造坐标:
import numpy as np arr = np.arange(16).reshape(4, 4) k_values = [-1, 0, 1] diags = [np.diag(arr, k=k) for k in k_values]
但注意:这些结果长度不同(例如 k=-1 和 k=1 在非方阵中可能差 1),直接 np.stack 会失败。真要对齐,得先算每条长度:min(m + k, n - k)(对 k >= 0),再统一截断或补零。
立即学习“Python免费学习笔记(深入)”;
性能敏感场景慎用 np.diag,主对角线可用视图替代
如果只取主对角线(k=0)且后续还要修改原矩阵对应位置,np.diag 返回的是副本,改了没用。此时应直接用 np.arange 配合花式索引获取视图:
diag_view = arr[np.arange(min(arr.shape)), np.arange(min(arr.shape))]
这个 diag_view 是原数组的视图,修改它等于修改原矩阵对角线元素。而 np.diag(arr) 总是拷贝。另外,在循环中高频调用 np.diag(尤其小矩阵)会有明显函数调用开销,预计算索引数组更快。
偏移对角线无法用纯视图实现,必须拷贝,这是 np.diag 不可绕过的成本。


















