如何用numpy向量化运算将三维数组扩展为对角矩阵数组？

首先明确核心需求：对每个(n,m)，B[n,m]是一个D×D矩阵，其左上角A[n,m]×A[n,m]区域为单位矩阵，其余位置为0（假设A[n,m] ≤ D；若A是(N,M,D)形状，且每个(n,m)对应的D个元素均为同一标量，可先通过A = A[...,0]压缩为(N,M)形状）。

方法1：广播掩码赋值 #

这是最直观的向量化方式，通过生成掩码定位需要设为1的对角线位置：

import numpy as np

# 示例输入
N, M, D = 2, 3, 4
A = np.random.randint(1, D+1, size=(N, M))  # 若A是(N,M,D)，则A = A[..., 0]

# 初始化目标数组
B = np.zeros((N, M, D, D), dtype=A.dtype)

# 生成对角线索引
diag_indices = np.arange(D)

# 广播生成掩码：判断每个对角线位置是否在A[n,m]的范围内
mask = diag_indices < A[..., None]

# 为对角线符合条件的位置赋值1
B[..., diag_indices, diag_indices] = mask.astype(A.dtype)

方法2：基于单位矩阵模板的广播截断 #

先创建一个D×D的单位矩阵模板，再通过广播掩码保留对应区域：

import numpy as np

N, M, D = 2, 3, 4
A = np.random.randint(1, D+1, size=(N, M))

# 生成D×D单位矩阵模板
eye_template = np.eye(D, dtype=A.dtype)

# 扩展A的维度以匹配模板的广播要求，生成掩码
mask = np.arange(D) < A[..., None, None]

# 掩码与模板相乘得到目标数组
B = eye_template * mask

特殊情况：若A[n,m]恒等于D #

如果所有A[n,m]都是D，那直接广播单位矩阵即可：

B = np.tile(np.eye(D), (N, M, 1, 1))

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