如何优化大型pandas DataFrame迭代子集均值及比率计算效率

原实现的核心性能瓶颈在于逐列循环+query()调用存在极高的Python层开销，3万次循环的重复语法解析、子集切片操作会浪费大量时间，我们可以通过全向量化操作消除Python循环，把所有计算转到C层面执行。

核心实现代码 #

import pandas as pd
import numpy as np

# 1. 对齐阈值和DataFrame的列顺序
high_cut = pd.Series(cutoff_vals['high']).reindex(df.columns)
low_cut = pd.Series(cutoff_vals['low']).reindex(df.columns)

# 2. 一次性生成全量高低阈值布尔掩码（和df同形状）
mask_high = df > high_cut
mask_low = df < low_cut

# 3. 转浮点型用于矩阵运算，float32可节省一半内存，精度不足可换float64
m_h = mask_high.astype('float32')
m_l = mask_low.astype('float32')

# 4. 矩阵乘法批量计算所有子集的列总和
# 输出形状为(N,N)：行对应筛选依据的原列序号，列对应计算均值的列序号
sum_high = m_h.T @ df
sum_low = m_l.T @ df

# 5. 除以各子集的有效行数得到均值矩阵，自动处理除0场景填充NaN
count_h = m_h.sum(axis=0)
count_l = m_l.sum(axis=0)
mean_high = sum_high.div(count_h, axis=0)
mean_low = sum_low.div(count_l, axis=0)

# 6. 直接逐元素计算得到最终N×N比率矩阵
ratios_matrix = mean_high / mean_low

额外优化说明 #

• 如果需要进一步提速，可以把数据转成Numpy数组执行运算，最后再转回带列名的DataFrame即可：

df_arr = df.values
m_h_arr = m_h.values
m_l_arr = m_l.values

sum_high_arr = m_h_arr.T @ df_arr
sum_low_arr = m_l_arr.T @ df_arr

mean_high_arr = sum_high_arr / count_h.values.reshape(-1,1)
mean_low_arr = sum_low_arr / count_l.values.reshape(-1,1)

ratios_matrix = pd.DataFrame(mean_high_arr / mean_low_arr, index=df.columns, columns=df.columns)

• 若某列的高低阈值筛选后无有效行，结果会自动填充NaN，可根据业务需求用fillna()自定义填充规则
• 运算前建议先校验high_cut、low_cut的索引和df.columns完全匹配，避免索引错位导致计算错误

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