如何在Python中使用NumPy对多维数组进行对角线提取？-Python教程-PHP中文网

np.diag 只提取一维对角线是因为其设计定位为矩阵代数中的标准对角操作，仅支持1d或2d输入；对高维数组需用 np.diagonal 沿指定两轴提取。

如何在python中使用numpy对多维数组进行对角线提取？

np.diag 为什么只提取一维对角线？

np.diag 默认只处理一维或二维数组：传入一维数组时返回其构成的对角矩阵；传入二维数组时只提取主对角线（k=0）。它不支持直接从三维及以上数组中按“体对角线”或“面内对角线”批量提取——这不是 bug，而是设计定位明确：它面向矩阵代数中的标准对角操作。

常见错误现象：np.diag 对三维数组报 ValueError: Input must be 1- or 2-dimensional；或者你传入 shape=(3,4,5) 却期望得到 3 条长度为 4 的对角线，结果发现根本不可行。

若目标是提取某个二维切片（如第 0 个平面）的对角线：先用索引切出 a[0]，再调用 np.diag(a[0])
若想对每个二维切片都做同样操作，用列表推导式：[np.diag(a[i]) for i in range(a.shape[0])]，但注意这会丢失向量化优势
性能影响：纯 Python 循环 + np.diag 在大数组上明显慢于原生 NumPy 操作

用 np.diagonal 提取高维数组的指定轴对角线

np.diagonal 是真正支持多维的函数。它沿指定两个轴提取对角线，其余轴保持不变。例如三维数组 a 的 shape=(4,5,6)，np.diagonal(a, axis1=1, axis2=2) 将在每个 a[i]（即每个 5×6 矩阵）中提取主对角线，返回 shape=(4,5) 的数组（因 min(5,6)=5）。

axis1 和 axis2 必须不同，且都在数组维度范围内；默认是 axis1=0, axis2=1
offset 参数控制偏移量（如 k=1 提取上一条对角线），但仅对所选两轴生效
结果形状：非对角轴尺寸不变，对角轴尺寸变为 min(shape[axis1], shape[axis2])（除非指定 offset 导致截断）
兼容性注意：NumPy np.diagonal 对 offset≠0 的行为在某些高维组合下有边界差异，建议升级到 1.19+

手动构造索引提取任意“空间对角线”

当需要提取真正的体对角线（如三维中满足 i==j==k 的所有元素），np.diagonal 也无能为力——它只管两个轴。这时必须手动生成坐标索引：

python学习

Python 编程学习助手 skill，面向零基础用户。当用户表达以下意图时触发： - 学习 Python 基础概念（变量、数据类型、循环、函数、类等） - 请求 Python 练习题或测验 - 请求审查或调试 Python 代码 - 请求 Python 项目实战指导（爬虫、数据分析、Web 开发等入门项目） - 泛化表述如"教我 Python""我想学编程""Python 怎么写""帮我学 Python" 此 skill

下载

import numpy as np
a = np.arange(27).reshape(3, 3, 3)
# 提取 i==j==k 的体对角线（长度为 3）
diag_3d = a[np.arange(3), np.arange(3), np.arange(3)]

这种写法本质是高级索引，灵活但需确保各索引数组长度一致。

若要提取 i==j, k 自由变化的“面内对角线平面”，可广播：a[np.arange(3)[:, None], np.arange(3)[:, None], np.arange(3)] 得到 (3,3,3) 子集
容易踩的坑：用 range 替代 np.arange 会导致索引失效（返回标量而非数组）
内存提示：高级索引会触发副本，超大数组慎用；如只需视图，优先考虑 np.diagonal 或 np.take 配合 np.eye 逻辑