如何用向量填充3D numpy数组中每个矩阵的对角线

现有一个3D NumPy数组，其中每个2D切片代表一个独立矩阵。需要将每个矩阵的对角线元素替换为一组特定值。

例如，对于一个3x3x3的数组：

import numpy as np
arr = np.array([[[a1, a2, a3],
                 [a4, a5, a6],
                 [a7, a8, a9]],

                [[b1, b2, b3],
                 [b4, b5, b6],
                 [b7, b8, b9]],

                [[c1, c2, c3],
                 [c4, c5, c6],
                 [c7, c8, c9]]])

需要把每个子矩阵的对角线[a1, a5, a9]、[b1, b5, b9]和[c1, c5, c9]分别替换为对应的新值，该如何实现？

解决方案 #

方法1：向量化高级索引（高效推荐） #

假设你有一个形状为(n, m, m)的3D数组，以及一个对应新对角线值的(n, m)数组new_diags，可以通过索引直接定位所有对角线元素并赋值：

import numpy as np

# 示例原数组
arr = np.array([[[1,2,3], [4,5,6], [7,8,9]],
                [[10,11,12], [13,14,15], [16,17,18]],
                [[19,20,21], [22,23,24], [25,26,27]]])

# 新对角线值：每行对应一个子矩阵的对角线
new_diags = np.array([[0, 0, 0], [1,1,1], [2,2,2]])

# 生成索引：i对应第几个矩阵，j对应对角线的行/列位置
i = np.arange(arr.shape[0])[:, None]
j = np.arange(arr.shape[1])

# 批量赋值
arr[i, j, j] = new_diags

这种方法完全利用NumPy的向量化操作，避免循环，处理大数组时效率远超循环。

方法2：循环逐个处理（直观易懂） #

如果更倾向于直观的逻辑，可以遍历每个2D子矩阵，用np.fill_diagonal替换对角线：

import numpy as np

arr = np.array([[[1,2,3], [4,5,6], [7,8,9]],
                [[10,11,12], [13,14,15], [16,17,18]],
                [[19,20,21], [22,23,24], [25,26,27]]])

new_diags = [[0,0,0], [1,1,1], [2,2,2]]

for mat_idx in range(arr.shape[0]):
    np.fill_diagonal(arr[mat_idx], new_diags[mat_idx])

这种写法逻辑清晰，适合刚接触NumPy的开发者，小尺寸数组下性能差异可以忽略。

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