如何在不使用嵌套循环的情况下初始化包含单位矩阵的多维NumPy数组?
当然有啦!NumPy的广播机制和数组操作工具就是专门用来解决这种要避免循环的场景的,给你几种简单好用的方法:
方法一:利用广播直接赋值(最简洁)
这是我最推荐的方式,先初始化一个全零的目标数组,然后直接把单位矩阵赋值给整个数组就行。NumPy的广播规则会自动帮你把单位矩阵“扩展”到前面的所有计数维度上,完全不用写任何循环。
对应你第一个例子(形状(3,2,2)):import numpy as np A = np.zeros((3, 2, 2)) A[...] = np.eye(2)第二个例子(形状
(3,4,5,5)):B = np.zeros((3,4,5,5)) B[...] = np.eye(5)是不是超省心?广播会自动处理前面的所有维度,把单位矩阵复制到每个对应的位置。
方法二:用
np.tile直接构造数组
如果你不想先初始化全零数组,也可以直接用np.tile来复制单位矩阵到指定维度。首先给单位矩阵加几个前置的空维度(用[None, ...]就能快速实现),然后指定每个计数维度上要复制的次数:
第一个例子的写法:A = np.tile(np.eye(2)[None, ...], (3, 1, 1))这里
np.eye(2)[None, ...]把原本(2,2)的单位矩阵变成(1,2,2),tile的参数(3,1,1)表示在第一个维度复制3次,后面两个维度保持不变,最终就得到了每个位置都是单位矩阵的(3,2,2)数组。
第二个例子对应的写法:B = np.tile(np.eye(5)[None, None, ...], (3,4,1,1))这次给单位矩阵加了两个前置维度变成
(1,1,5,5),然后在第一、第二个维度分别复制3次和4次,就得到了你要的目标数组。方法三:用
np.broadcast_to创建只读视图(内存友好)
如果你的数组只是用来读取,不需要修改的话,可以用np.broadcast_to创建一个广播视图,这样不会额外占用内存(因为它只是原单位矩阵的视图,没有实际复制数据):
第一个例子的写法:A = np.broadcast_to(np.eye(2)[None, ...], (3,2,2))提醒一下,这个数组是只读的,如果之后需要修改里面的值,记得用
A.copy()把它转换成可写的数组。
备注:内容来源于stack exchange,提问作者Triceratops

