如何用Numpy加速矩阵指定行的均值计算？性能优化问询

1. 全选场景的快速优化 #

当selected_rows包含所有行（不管顺序是否打乱），直接用原数组计算均值，彻底避免副本创建的开销：

import numpy as np, time

x = np.random.random((100000, 500))
selected_rows = np.arange(100000)

t0 = time.time()
if len(selected_rows) == x.shape[0]:
    # 可选：检查索引是否无重复（确保是真正全选）
    count = np.bincount(selected_rows, minlength=x.shape[0])
    if np.all(count == 1):
        y = x.mean(axis=0)
    else:
        # 存在重复索引，按实际选中行计算
        y = x[selected_rows].mean(axis=0)
else:
    y = x[selected_rows].mean(axis=0)
print(time.time() - t0)  # 速度接近原生mean的36ms

如果能确定selected_rows不会有重复索引，可简化判断：

if len(selected_rows) == x.shape[0]:
    y = x.mean(axis=0)
else:
    y = x[selected_rows].mean(axis=0)

2. 非全选场景的无副本计算 #

通过加权求和或矩阵乘法的方式，避免创建选中行的副本，提升计算效率：

方法A：矩阵乘法实现加权均值 #

t0 = time.time()
weights = np.zeros(x.shape[0], dtype=np.float64)
weights[selected_rows] = 1.0 / len(selected_rows)
y = x.T @ weights
print(time.time() - t0)

方法B：手动求和后除法 #

和mean逻辑一致，但可跳过mean内部的额外判断：

t0 = time.time()
y = x[selected_rows].sum(axis=0) / len(selected_rows)
print(time.time() - t0)

3. 用np.take替代切片 #

在部分场景下，np.take比直接切片的内存效率更高：

t0 = time.time()
y = np.take(x, selected_rows, axis=0).mean(axis=0)
print(time.time() - t0)

直接调用x.mean(axis=0)时，NumPy在连续内存块上计算，缓存命中率高，速度快。而用整数索引x[selected_rows, :]时，NumPy会创建副本（而非视图），新数组内存不连续，缓存效率骤降——哪怕索引是全行的顺序排列，副本创建的额外开销依然存在。

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