如何在NumPy中高效实现块常量矩阵与矩阵的加法？

问题背景 #

我们有：

• n×n矩阵A
• m×m矩阵B
• 块大小向量b=[b₁,...,bₘ]，满足b₁+…+bₘ=n且bᵢ≥1

定义n×n块矩阵B_block，其(i,j)块为值为Bᵢⱼ的bᵢ×bⱼ常量矩阵。需要在NumPy中高效计算A+B_block。

当前实现可行但效率较低，代码如下：

import numpy as np

A = np.arange(36).reshape(6,6)
B = np.arange(9).reshape(3,3)
blocks_sizes = np.array([3,2,1])
B_block = np.block([[np.ones((a,b))*B[i,j] for j,b in enumerate(block_sizes)] for i,a in enumerate(block_sizes)])
C = A + B_block
print(A)
print(B_block)
print(C)

运行结果：

[[ 0  1  2  3  4  5]
 [ 6  7  8  9 10 11]
 [12 13 14 15 16 17]
 [18 19 20 21 22 23]
 [24 25 26 27 28 29]
 [30 31 32 33 34 35]]
[[0. 0. 0. 1. 1. 2.]
 [0. 0. 0. 1. 1. 2.]
 [0. 0. 0. 1. 1. 2.]
 [3. 3. 3. 4. 4. 5.]
 [3. 3. 3. 4. 4. 5.]
 [6. 6. 6. 7. 7. 8.]]
[[ 0.  1.  2.  4.  5.  7.]
 [ 6.  7.  8. 10. 11. 13.]
 [12. 13. 14. 16. 17. 19.]
 [21. 22. 23. 25. 26. 28.]
 [27. 28. 29. 31. 32. 34.]
 [36. 37. 38. 40. 41. 43.]]

高效解决方案 #

可以利用NumPy的np.repeat函数实现完全向量化的操作，无需Python循环或手动构造块矩阵：

import numpy as np

A = np.arange(36).reshape(6,6)
B = np.arange(9).reshape(3,3)
blocks_sizes = np.array([3,2,1])

# 先按行重复对应块大小，再按列重复
B_block = np.repeat(np.repeat(B, blocks_sizes, axis=0), blocks_sizes, axis=1)
C = A + B_block

print(A)
print(B_block)
print(C)

原理说明 #

• np.repeat(B, blocks_sizes, axis=0)：将B的每一行按blocks_sizes中的对应次数重复，比如B的第0行重复3次、第1行重复2次、第2行重复1次，得到6×3的中间矩阵。
• 再对这个中间矩阵按列方向用blocks_sizes重复，得到6×6的B_block，完全符合需求。
• 这种方式依赖NumPy底层优化的C实现，处理大矩阵时性能远高于手动构造块的方法，同时兼容块大小不一致的场景，是np.kron方法的灵活扩展。

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