我可以创建一个包含没有嵌套循环的单位矩阵的多维数组吗？

Question

假设我有一个 Numpy 数组

n

索引，其中第一个

n-2

代表一些计数索引，最后 2 个索引代表一个方

MxM

矩阵。我想初始化这个结构，以便它将包含单位矩阵的副本。

示例（此处 N=3，M=2）：

A = numpy.zeros((3,2,2))
for k in range(3):
     A[k,:,:] = numpy.eye(2)

另一个例子（这里 N=4，M=5）：

B = numpy.zeros((3,4,5,5))
for k1 in range(3):
     for k2 in range(4):
         B[k1,k2,:,:] = numpy.eye(5)

有没有办法在不依赖嵌套循环的情况下做到这一点？

Answer 1

你可以

repeat

:

A = np.repeat(np.eye(2)[None], 3, axis=0)

对于更复杂的情况，结合

reshape

：

extra = (3, 4)
M = 5
B = np.repeat(np.eye(M)[None], np.prod(extra), axis=0).reshape(extra+(M, M))

或与

tile

:

extra = (3, 4)
B = np.tile(np.eye(5), extra+(1, 1))

或者，从

numpy.zeros

使用索引：

B = np.zeros((3, 4, 5, 5))
x = np.arange(B.shape[-1])
B[..., x, x] = 1

Answer 2

“紧凑”并不是一个非常具体的术语，因此很难说哪种方法可以满足要求。然而，一种方法是按顺序使用

numpy.tile

、

numpy.split

和

numpy.stack

。使用

tile

排列所有单位矩阵，使用

split

将它们分解为列表，然后使用

stack

将它们放入单个多维数组中。这真的比循环更紧凑吗？谁知道呢，但你可以用一行来完成。

您还可以使用列表理解，它不会消除循环，而是将其放在一行中。