如何在Matlab中从三维数组提取对角切片并存储为新三维数组

适用场景 #

针对尺寸为m×n×k的三维矩阵，提取所有行号等于列号的对角切片，最终输出为三维数组。

实现代码 #

方法1：循环实现（推荐大矩阵使用） #

逻辑清晰，内存占用可控，Matlab新版本对循环优化到位，运行效率无明显短板。

% 输入三维矩阵A
[m, n, k] = size(A);
diag_length = min(m, n);
% 初始化结果数组，可根据需求调整维度顺序
result = zeros(diag_length, k, 1);

for idx = 1:diag_length
    % 提取第idx组对角元素，squeeze用于去除冗余维度
    result(idx, :, 1) = squeeze(A(idx, idx, :));
end

方法2：向量化实现（推荐中小矩阵使用） #

代码简洁，无需写循环，可读性更高。

% 输入三维矩阵A
[m, n, k] = size(A);
diag_length = min(m, n);
% 生成对角线元素的线性索引
diag_idx = sub2ind([m, n], 1:diag_length, 1:diag_length);
% 提取元素并重排为三维数组
result = reshape(A(diag_idx, :), diag_length, k, 1);

补充说明 #

• 如果需要提取反对角切片，仅需将索引中的列号替换为m+1-idx即可
• 若需要其他维度组合的对角切片（如列号等于页号、行号等于页号），调整索引的维度对应关系即可
• 可根据后续使用需求，修改reshape的参数调整输出数组的维度顺序

内容的提问来源于stack exchange，提问作者user3482383