如何在NumPy中根据输入维度生成对应单位矩阵数组

需求说明：
编写一个函数，接收NumPy数组作为参数：

• 若输入是shape=(3,)的一维数组（表示单个3D物理向量），则返回3×3单位矩阵
• 若输入是shape=(3,n)的二维数组（表示n个3D物理向量），则返回shape=(3,3,n)的三维数组，其中每个unit[:,:,k]都是3×3单位矩阵
  要求实现高效，避免手动逐个赋值对角元素的方式。

高效实现方案 #

可以利用NumPy的广播机制或内置函数的维度扩展能力，两种高效实现方式如下：

方式1：利用np.eye维度扩展 + 广播（无数据复制，性能最优）

通过给单位矩阵添加新维度，再用np.broadcast_to扩展到目标形状，这种方式不会额外复制数据，是性能最高的实现：

import numpy as np

def func(x):
    if x.ndim == 1:
        unit = np.identity(3)
    elif x.ndim == 2:
        n = x.shape[1]
        # 先创建(3,3,1)的单位矩阵，再广播到(3,3,n)
        unit = np.broadcast_to(np.eye(3)[..., None], (3,3,n))
    return unit

方式2：利用np.tile进行维度复制（代码更简洁）

如果不需要严格的无复制视图，也可以用np.tile快速复制单位矩阵的维度：

import numpy as np

def func(x):
    if x.ndim == 1:
        unit = np.identity(3)
    elif x.ndim == 2:
        n = x.shape[1]
        # 将(3,3)的单位矩阵扩展为(3,3,1)后，复制n次
        unit = np.tile(np.eye(3)[..., None], (1,1,n))
    return unit

说明 #

• 两种方式都避免了手动赋值对角元素的繁琐操作，且都是NumPy底层优化的实现，性能远高于手动循环或赋值
• np.broadcast_to返回原数组的视图，不占用额外内存；np.tile会创建新数组副本，内存占用稍高，但代码更直观

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