Pandas快速计算指定窗口内列值之和的优化方案咨询

看起来你在处理基因组数据的窗口求和问题，55k行数据运行慢确实头疼——主要原因是你的当前实现每次都要全表扫描做布尔筛选，没有利用数据的有序性和分组特性。下面给你几个逐步优化的方案，从简单调整到极致高效：

1. 先按染色体分组，避免全表扫描 #

如果你的数据是按chrom和pos递增排列的（基因组数据通常都是这样），可以提前按chrom分组，后续只在目标染色体的子集中操作，大幅减少每次处理的数据量：

# 提前执行一次分组，重复调用函数时复用这个分组对象
chrom_groups = vdf.groupby('chrom')

def mafWindow_optimized1(chrom_groups, polyChrom, polyPos, distance, windowSize):
    # 先获取目标染色体的子集，跳过全表筛选
    try:
        df_chrom = chrom_groups.get_group(polyChrom)
    except KeyError:
        return 0.0  # 无对应染色体数据时返回0
    
    # 计算两个窗口的起止边界，处理pos为负数的情况
    start1 = max(polyPos - distance, 0)
    end1 = start1 + windowSize
    start2 = polyPos + distance - windowSize
    end2 = polyPos + distance
    
    # 仅在目标染色体子集中筛选位置范围
    mask = ((df_chrom['pos'] >= start1) & (df_chrom['pos'] <= end1)) | \
           ((df_chrom['pos'] >= start2) & (df_chrom['pos'] <= end2))
    return df_chrom.loc[mask, 'AFR'].sum()

2. 二分查找定位窗口（最快方案） #

如果每个染色体组内的pos是严格递增有序的，用bisect模块做二分查找可以直接定位窗口的起止索引，完全避免逐行判断，时间复杂度降到O(logN)：

import bisect

# 提前预处理：按chrom分组，每个组保存排序后的pos数组和AFR数组
# 如果原数据已经按pos排序，可以跳过sort_values步骤
preprocessed_data = {}
for chrom, group in vdf.groupby('chrom'):
    sorted_group = group.sort_values('pos').reset_index(drop=True)
    preprocessed_data[chrom] = {
        'pos_array': sorted_group['pos'].values,
        'afr_array': sorted_group['AFR'].values
    }

def mafWindow_optimized2(preprocessed_data, polyChrom, polyPos, distance, windowSize):
    chrom_data = preprocessed_data.get(polyChrom)
    if not chrom_data:
        return 0.0
    
    pos_arr = chrom_data['pos_array']
    afr_arr = chrom_data['afr_array']
    
    # 计算两个窗口的边界
    start1 = max(polyPos - distance, 0)
    end1 = start1 + windowSize
    start2 = polyPos + distance - windowSize
    end2 = polyPos + distance
    
    # 用二分查找快速定位区间索引
    left_idx1 = bisect.bisect_left(pos_arr, start1)
    right_idx1 = bisect.bisect_right(pos_arr, end1)
    left_idx2 = bisect.bisect_left(pos_arr, start2)
    right_idx2 = bisect.bisect_right(pos_arr, end2)
    
    # 计算两个区间的AFR和（区间无数据时返回0）
    sum_window1 = afr_arr[left_idx1:right_idx1].sum() if right_idx1 > left_idx1 else 0
    sum_window2 = afr_arr[left_idx2:right_idx2].sum() if right_idx2 > left_idx2 else 0
    
    return sum_window1 + sum_window2

这个方案的优势在于预处理只做一次，后续每次查询都是毫秒级的，哪怕你要处理上万个polyPos也能轻松应对。

3. 额外小技巧：用NumPy替代Pandas操作 #

上面的预处理已经用到了NumPy数组，因为NumPy的切片求和比Pandas的loc+sum更快，避免了Pandas的索引和对象开销，进一步提升效率。

原代码慢的核心原因 #

你的原函数每次调用都要对整个DataFrame做布尔索引筛选：vdf.loc[(vdf['chrom'] == polyChrom) & ...]，这会遍历全表的每一行。55k行单次调用可能还好，但如果是循环处理多个目标位置，累计的时间成本就会非常高。优化后的方案要么缩小了处理范围，要么用二分查找直接定位，彻底避免了全表扫描。

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