Numpy矩阵高效聚合与值复制的实现方法

一、矩阵聚合（N×N → P×P） #

针对多对一的维度映射关系，推荐两种Numpy原生的高效向量化实现，彻底替代低效的嵌套循环：

准备工作：将字典映射转为索引数组 #

先把原维度到新维度的字典映射转换为Numpy数组，这是后续所有操作的基础：

import numpy as np

# 示例矩阵
m = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
# 原维度到新维度的映射字典
big2small = {0:0, 1:1, 2:0, 3:1}
# 转换为索引数组：索引对应原维度下标，值对应新维度下标
idx = np.array([big2small[i] for i in range(m.shape[0])])
# idx = array([0, 1, 0, 1])

方法1：使用np.add.at直接累加（最直观的向量化累加） #

利用np.add.at实现对目标矩阵的批量累加操作，完全规避循环：

# 新矩阵的尺寸P
p = idx.max() + 1
# 初始化结果矩阵
result = np.zeros((p, p), dtype=m.dtype)
# 按映射关系批量累加原矩阵的值到结果矩阵
np.add.at(result, (idx[:, None], idx), m)

print(result)
# 输出：
# [[24 28]
#  [40 44]]

该方法直接操作底层数组，效率极高，适配1000×1000级别的大矩阵。

方法2：利用矩阵乘法（基于One-Hot编码） #

通过将索引转换为One-Hot矩阵，利用矩阵乘法实现分组求和，数学逻辑清晰：

p = idx.max() + 1
# 生成One-Hot矩阵，shape=(N, P)
one_hot = np.eye(p)[idx]
# 矩阵乘法实现行、列的分组求和
result = one_hot.T @ m @ one_hot

print(result)
# 输出与方法1一致

矩阵乘法依托Numpy的BLAS优化，性能与np.add.at相当，适合偏好数学表达的场景。

二、矩阵扩展（P×P → N×N） #

反向操作时，直接利用索引数组的广播特性即可实现值的批量复制，完全向量化：

# 示例小矩阵
small_m = np.array([[1,2],[3,4]])
# 利用索引数组的广播索引，直接生成扩展矩阵
expanded_m = small_m[idx[:, None], idx]

print(expanded_m)
# 输出：
# [[1 2 1 2]
#  [3 4 3 4]
#  [1 2 1 2]
#  [3 4 3 4]]

该操作完全依赖Numpy的索引优化，无任何循环，是效率最高的实现方式。

注意事项 #

如果原映射字典的新维度下标不是连续整数（比如新索引为[2,3,2,3]），可先通过np.unique转换为连续下标，处理完成后再映射回原索引：

# 示例非连续映射
big2small = {0:2, 1:3, 2:2, 3:3}
values = list(big2small.values())
# 转换为连续下标
unique_new_idx, idx = np.unique(values, return_inverse=True)
# 此时idx = array([0,1,0,1])，后续操作同上述方案
# 处理完成后，可通过unique_new_idx将结果的下标映射回原索引

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