如何将多个二维方阵沿对角线拼接为更大尺寸的二维方阵

解决方案 #

你需要实现的操作是构建块对角矩阵，NumPy 已经提供了高度优化的内置实现，无需自行编写复杂的循环判断逻辑。

方案1：直接调用内置函数（推荐） #

NumPy 的np.linalg.block_diag方法原生支持将任意数量的二维方阵拼接为块对角矩阵，一行代码即可完成操作：

import numpy as np

# ind_array 为存储所有子方阵的列表
coll_array = np.linalg.block_diag(*ind_array)

其中*操作符会自动解包列表中的所有子方阵作为函数入参，无论你有多少个待拼接的子方阵都可以直接使用，不需要额外修改代码。

方案2：手动实现高效版本（如需自定义逻辑） #

如果你需要了解底层实现逻辑，或者要在拼接过程中加入自定义操作，可以通过切片赋值的方式实现，完全避免逐元素判断，效率和内置函数接近：

import numpy as np

# 计算每个子方阵的阶数
sub_sizes = [arr.shape[0] for arr in ind_array]
# 计算大矩阵总阶数
total_size = sum(sub_sizes)
# 初始化全0大矩阵，dtype和子矩阵保持一致
coll_array = np.zeros((total_size, total_size), dtype=ind_array[0].dtype)

current_offset = 0
for sub_arr in ind_array:
    n = sub_arr.shape[0]
    # 直接给对角块对应的切片赋值
    coll_array[current_offset:current_offset+n, current_offset:current_offset+n] = sub_arr
    current_offset += n

方案优势 #

• 适配任意数量的子方阵，不需要手动编写任何条件判断语句
• 基于NumPy向量化操作实现，效率远高于逐坐标遍历的实现，处理大规模数据时速度提升可达数百倍
• 兼容子方阵内部元素各不相同的场景，适用范围更广

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