高效构建由对称子矩阵组成的三维数组

本文介绍如何利用 numpy 的广播式高级索引，无需显式循环即可批量提取二维数组中多个由不同索引集定义的对称子矩阵，并堆叠为三维数组，显著提升性能。 本文介绍如何利用 numpy 的广播式高级索引，无需显式循环即可批量提取二维数组中多个由不同索引集定义的对称子矩阵，并堆叠为三维数组，显著提升性能。在科学计算与矩阵分析中，常需从一个大型二维数组 A（如协方差矩阵、相似度矩阵）中，按多组索引集合批量提取对称子矩阵（即形如 A[idx, :][:, idx] 的方阵），并将结果统一组织为三维数组 B，其中 B[i] 对应第 i 组索引所确定的子矩阵。若采用 Python 循环配合列表推导式（如 [A[idx, :][:, idx] for idx in I]），虽逻辑清晰，但在大规模数据下性能极差——每次索引均触发独立内存拷贝与 Python 解释器开销。幸运的是，NumPy 提供了广播式高级索引（broadcasted advanced indexing） 这一强大机制，可一次性完成全部子矩阵提取。核心技巧在于合理扩展索引数组 I 的维度，使其与 A 的二维结构对齐：B = A[I[:, :, None], I[:, None, :]]该表达式中：I 形状为 (N, K)，表示 N 组、每组 K 个索引；I[:, :, None] 将其扩展为 (N, K, 1)，用于索引 A 的行维度；I[:, None, :] 将其扩展为 (N, 1, K)，用于索引 A 的列维度；利用 NumPy 广播规则，二者共同生成 (N, K, K) 的索引组合，最终 A[...] 返回形状为 (N, K, K) 的三维数组 B，完美对应需求。? 验证示例： 腾讯小微 基于微信AI智能对话系统打造的智能语音助手解决方案