如何以向量化方式实现numpy数组的行分组均值计算？

问题背景 #

给定分组索引列表：

outer_list = [[0], [1, 2], [1, 2], [3, 5], [4], [3, 5]]

其中每个子列表对应父列表中该索引所属的分组（例如索引1和2同属一组，索引3和5同属一组）。同时有一个numpy矩阵val_mat：

val_mat = np.arange(2*6).reshape(2, 6)

需求是对每行数据，将同一分组内的元素替换为该分组的行均值，最终得到均值矩阵。

当前实现通过遍历outer_list计算均值，在处理10k×500k规模矩阵时耗时约2分钟，需要更高效的向量化实现。

向量化实现步骤 #

1. 生成分组标签数组 #

首先将重复的分组映射为唯一ID，生成每个索引对应的分组标签：

import numpy as np

outer_list = [[0], [1, 2], [1, 2], [3, 5], [4], [3, 5]]
# 将子列表转为元组（可哈希），用于去重
group_tuples = [tuple(lst) for lst in outer_list]
# 获取唯一分组并生成每个索引的分组ID
unique_groups, group_ids = np.unique(group_tuples, return_inverse=True)

此时group_ids为array([0, 1, 1, 2, 3, 2])，代表每个索引对应的分组编号。

2. 向量化计算分组均值 #

利用np.bincount实现高效的分组求和与计数，进而计算均值并映射回原索引位置：

val_mat = np.arange(2*6).reshape(2, 6)

# 计算每个分组的元素个数
group_counts = np.bincount(group_ids)

# 对每行计算分组均值并生成结果矩阵
mean_vals = np.zeros_like(val_mat)
for i in range(val_mat.shape[0]):
    row = val_mat[i]
    # 分组求和
    group_sums = np.bincount(group_ids, weights=row)
    # 分组均值
    group_means = group_sums / group_counts
    # 将均值映射到对应索引位置
    mean_vals[i] = group_means[group_ids]

3. 进一步优化：避免行循环（可选） #

如果想彻底去掉行循环，可以利用广播和矩阵运算实现：

# 构建分组计数的广播矩阵：(n_groups, n_cols)
count_matrix = group_counts[group_ids][np.newaxis, :]
# 构建分组求和矩阵：(n_rows, n_groups)
sum_matrix = np.dot(val_mat, np.eye(len(group_counts))[group_ids])
# 计算均值矩阵
mean_vals = sum_matrix / count_matrix

注：该方法在分组数量极大时（如50万组），np.eye会生成超大矩阵，内存可能不足，此时优先选择带行循环的版本（行循环的开销远小于原方法的50万次切片操作）。

效果对比 #

原方法需要遍历50万次分组并执行切片运算，向量化方法仅需几次numpy内置操作，处理10k×500k矩阵的耗时可大幅降低（通常能压缩到数十秒甚至更短）。

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