Numpy中行权重加权求和的高效实现方法

这事儿其实用Numpy的矩阵乘法就能轻松搞定，而且是性能最优的方案——毕竟Numpy依赖的BLAS/LAPACK库都是经过高度优化的底层实现，比你手动写Python循环快得多。

核心思路 #

你要的是对每个k，用gamma[:,k]（N维向量）作为权重，对X的每一行（D维向量）做加权求和，最终得到K个D维向量（也就是形状为(K,D)的数组）。

这里的关键是利用矩阵乘法的维度匹配：

• X的形状是(N,D)
• gamma转置后形状是(K,N)
• 两者相乘gamma.T @ X会得到(K,D)的结果，其中第k行正好就是你要的s[k]。

代码示例 #

import numpy as np

# 生成测试用的随机数组
N, D, K = 100, 20, 15
X = np.random.randn(N, D)
gamma = np.random.randn(N, K)

# 高效计算加权和
s = gamma.T @ X  # 或者用 np.dot(gamma.T, X)，效果完全一致

# 验证结果正确性（和手动计算对比）
manual_s0 = np.sum(gamma[:,0, np.newaxis] * X, axis=0)
print(np.allclose(s[0], manual_s0))  # 输出True，说明结果完全匹配

为什么这是最优解？ #

• 向量化操作：完全避免了Python层面的循环，所有计算都在底层的编译代码中执行，对于大尺寸数组（比如N、K、D都上千），速度能提升几十甚至上百倍。
• 内存高效：矩阵乘法的实现会尽可能复用内存，减少不必要的中间数组创建。

反面例子：别这么写！ #

如果你手动写循环来计算，性能会差很多，比如：

# 不推荐的循环写法，仅作性能对比
s_loop = np.zeros((K, D))
for k in range(K):
    s_loop[k] = np.sum(gamma[:,k].reshape(-1,1) * X, axis=0)

用%timeit测试的话，这个循环版本的耗时会是矩阵乘法版本的几十倍，数据量越大差距越明显。

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