如何不使用for循环批量生成2×2矩阵并存入numpy ndarray

架构老林

2026-9-28

Numpy批量生成对角矩阵的高效解决方案

核心思路

利用Numpy的广播机制实现完全矢量化的批量生成，全程无显式循环，计算效率远高于append或循环实现。

实现代码

import numpy as np
x = np.arange(1, 4, 1)

# 直接通过广播生成形状为 (N, 2, 2) 的目标数组
y3 = x[:, np.newaxis, np.newaxis] * np.eye(2)

# 若需要整数类型，可添加类型转换
# y3 = y3.astype(np.int32)

结果验证

生成的y3形状为(3, 2, 2)，输出如下，完全符合需求：

array([[[1., 0.],
        [0., 1.]],

       [[2., 0.],
        [0., 2.]],

       [[3., 0.],
        [0., 3.]]])

原代码告警原因说明

之前实现的example2中，当传入的d是一维数组时，[[d,0],[0,d]]的子元素分别是长度为N的数组和标量，形状不匹配，Numpy无法将其拼接为规整的高维数组，因此会触发VisibleDeprecationWarning警告，同时生成元素为数组的object类型数组，不符合预期结构。