如何优化Python中的enumerate循环？实现纯矩阵运算提速

问题背景 #

原代码通过enumerate循环处理5000个随机样本，运行效率极低，需要重构为纯向量化/矩阵运算来提速。

输入参数 #

import numpy as np
from scipy.stats import rv_discrete

T  =  3
N  =  5000
k  =  [0.1193354,0.27447143,0.42960745]
P  =  [0.31600902,0.02144897,0.03998757,0.00708576,0.27952587,0.02057329,0.01790612,0.00689641,0.29056699]
cont  =  np.arange(T*T)
G     =  np.zeros([T,T])
# 注：需预先定义T×T矩阵c和常数beta

原循环代码 #

dis   =  rv_discrete(values=(cont, P))

for t, i_t in enumerate(dis.rvs(size = N)):
    Q  =  np.zeros([T,T])
    if i_t // T == i_t % T:
        Q[i_t // T, i_t % T] = c[i_t // T, i_t % T] / (1 - beta)
    else:
        Q[i_t // T, i_t % T] = c[i_t // T, i_t % T]
    
    G  += Q / P[i_t]

核心逻辑 #

随机抽取N个样本，每个样本对应T×T矩阵Q的一个位置：

• 若样本索引i_t的行号（i_t//T）等于列号（i_t%T），则该位置取值为c[行,列]/(1-beta)
• 否则取值为c[行,列]
• 最终G是所有Q矩阵除以对应P[i_t]后的总和

解决方案：全向量化实现 #

完全移除循环，利用NumPy的向量化操作一次性计算所有样本的贡献，大幅提升运行速度：

import numpy as np
from scipy.stats import rv_discrete

# 初始化参数
T = 3
N = 5000
P = np.array([0.31600902,0.02144897,0.03998757,0.00708576,0.27952587,0.02057329,0.01790612,0.00689641,0.29056699])
cont = np.arange(T*T)
c = np.random.rand(T, T)  # 替换为你的实际c矩阵
beta = 0.5  # 替换为你的实际beta值

# 1. 一次性生成所有样本
dis = rv_discrete(values=(cont, P))
samples = dis.rvs(size=N)

# 2. 批量计算所有样本的行、列索引
rows = samples // T
cols = samples % T

# 3. 批量生成系数：对角线位置用1/(1-beta)，其余用1
coeffs = np.where(rows == cols, 1/(1 - beta), 1)

# 4. 获取对应c值并乘以系数
c_vals = c[rows, cols] * coeffs

# 5. 计算每个样本的权重（1/P[i_t]）
weights = 1 / P[samples]

# 6. 高效累加所有样本的贡献到G矩阵
G = np.zeros((T, T))
np.add.at(G, (rows, cols), c_vals * weights)

关键优化点 #

• 批量生成样本：避免循环中逐个生成样本，一次性获取所有数据
• 向量化索引计算：用NumPy数组直接批量计算所有样本的行、列位置
• 条件逻辑向量化：用np.where替代循环中的if-else，批量处理系数
• 高效累加：用np.add.at直接对矩阵指定位置做加权累加，替代循环中逐个创建Q矩阵再相加的低效操作

速度提升效果 #

原循环在N=5000时需反复创建小矩阵并逐次累加，向量化版本依赖NumPy的底层C实现，速度可提升10~100倍（具体幅度取决于T和N的规模）。

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