时间戳

首页 > 解决方案 > 用于在 NumPy 中创建块矩阵序列的向量化语法

python - 用于在 NumPy 中创建块矩阵序列的向量化语法

问题描述

我有两个 3D 数组A，分别B带有形状(k, n, n)和(k, m, m)。我想创建一个C形状矩阵，(k, n+m, n+m)使得对于每个0 <= i < k，二维矩阵C[i,:,:]是通过放置A[i, :, :]在左上n x n部分和B[i, :, :]右下m x m部分获得的块对角矩阵。

目前我正在使用以下方法来实现这一点，即 NumPy：

C = np.empty((k, n+m, n+m))
for i in range(k):
    C[i, ...] = np.block([[A[i,...], np.zeros((n,m))],
                          [np.zeros((m,n)), B[i,...]]])

我想知道是否有办法在没有for循环的情况下做到这一点。我认为如果k很大，我的解决方案不是很有效。

标签： pythonnumpy

解决方案

IIUC 您可以简单地切片和分配 -

C = np.zeros((k, n+m, n+m),dtype=np.result_type(A,B))
C[:,:n,:n] = A
C[:,n:,n:] = B

推荐阅读