如何用索引元组列快速提取DataFrame切片并计算均值？

问题场景 #

现有一个索引单调递增但非连续的DataFrame df，另有存储起始-结束索引元组的DataFrame idx_df，需要计算每个元组对应df切片的列均值。原用apply实现，但大数据量下速度极慢，需要高效的向量化方案。

原示例代码：

import pandas as pd
import numpy as np

df = pd.DataFrame({'col1': np.random.random(10), 'col2': np.random.random(10)}, index=sorted(np.random.choice(np.arange(30), size=10, replace=False)))
idx_df = pd.DataFrame({'idxs': [(0, 3), (2, 6), (3, 5), (6, 10)]})

# 原低效实现
idx_df.apply(lambda row: df.loc[row['idxs'][0]:row['idxs'][1], ['col1', 'col2']].mean(axis=0), axis=1)

向量化优化方案 #

核心思路是利用前缀和数组计算区间和，再结合区间元素个数得到均值，全程避免逐行循环处理：

1. 拆分索引元组：把idx_df里的起始、结束索引拆成单独列
2. 定位区间在df中的位置：用searchsorted找到每个起始/结束索引在df索引中的边界位置
3. 计算前缀和：生成包含初始0值的前缀和数组，方便快速计算任意区间的和
4. 计算区间和与均值：通过前缀和差值得到区间和，除以元素个数得到均值

具体代码：

# 1. 拆分起始、结束索引
idx_df[['start', 'end']] = pd.DataFrame(idx_df['idxs'].tolist(), index=idx_df.index)

# 2. 找到每个start/end在df索引中的位置
# left=True：找第一个>=start的位置；right=True：找第一个>end的位置（对应切片的右边界）
start_pos = df.index.searchsorted(idx_df['start'], side='left')
end_pos = df.index.searchsorted(idx_df['end'], side='right')

# 3. 计算前缀和（开头加一行0，方便计算从第一个元素开始的区间）
prefix_sum = pd.concat([pd.DataFrame([[0]*len(df.columns)], columns=df.columns), df.cumsum()], ignore_index=True)

# 4. 计算区间和与均值
interval_sum = prefix_sum.iloc[end_pos].values - prefix_sum.iloc[start_pos].values
counts = end_pos - start_pos
# 避免除以0的情况（如果区间内没有元素，均值设为NaN）
mean_df = pd.DataFrame(interval_sum / counts[:, np.newaxis], columns=df.columns)

性能优势 #

全程用Pandas/Numpy的向量化操作，彻底避开apply的逐行Python循环开销。前缀和计算是O(n)复杂度，后续区间查询都是O(1)，大数据量下性能提升可达数倍甚至数十倍。

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