Python中如何构造可指定填充对角线位置的3D对角矩阵？

云原生阿岚

2026-10-4

3D 自定义对角矩阵实现方案

你可以通过手动构造索引的方式实现支持指定偏移位置的3D对角矩阵填充，等效于2D下numpy.diag的3D扩展，实现逻辑如下：

核心实现代码

import numpy as np

def diag3d(arr_shape, offset=0, fill_value=1):
    """
    3D版本diag实现，支持自定义对角线偏移位置
    :param arr_shape: 目标3D数组形状，要求三个维度长度相同，格式为(N,N,N)
    :param offset: 对角线偏移量，正数向高索引偏移，负数向低索引偏移，0为主对角线
    :param fill_value: 对角线位置填充值
    :return: 填充完成的3D对角矩阵
    """
    N = arr_shape[0]
    # 维度合法性校验
    assert len(arr_shape) == 3 and arr_shape[1] == N and arr_shape[2] == N, "仅支持三个维度长度相同的3D矩阵"
    assert abs(offset) < N, "偏移量超出矩阵维度范围"
    
    res = np.zeros(arr_shape, dtype=np.array(fill_value).dtype)
    # 构造对应偏移的三维索引
    if offset >= 0:
        idx = np.arange(N - offset)
        res[idx, idx, idx + offset] = fill_value
    else:
        offset_abs = abs(offset)
        idx = np.arange(N - offset_abs)
        res[idx + offset_abs, idx, idx] = fill_value
    return res

使用示例

生成形状为(5,5,5)、偏移量为1、填充值为3的3D对角矩阵：
```
arr = diag3d((5,5,5), offset=1, fill_value=3)
```
如果需要在已有的3D数组上填充对角线，直接替换代码中初始化res的步骤为传入已有数组即可。
如果需要其他维度组合的对角线（比如仅前两个维度对角，第三个维度全局填充），只需调整索引赋值的维度顺序即可，例如res[idx + offset, idx, :] = fill_value即可实现前两个维度偏移的3D块对角矩阵。