如何高效统计anchor分组内指定取值区间的数值个数 避免内存溢出

内存溢出原因 #

你当前的实现采用了同anchor字段自连接的逻辑，会为每个分组生成 n² 行的中间结果（n为分组内的数据条数），当数据总大小达到3GB时，中间表的体积会膨胀到原始数据的数十甚至上百倍，直接触发内存溢出。

优化方案 #

采用排序+二分查找的方式替代自连接逻辑，无需生成大体积中间表，内存占用和原始数据集基本持平，时间复杂度也从O(n²)降低到O(n log n)，完全适配3GB级别的数据处理需求。

优化后代码 #

import pandas as pd

# 原有上下边界计算逻辑保留
df["rule_8_comp_low"] = df["y_val"] - df["anchor_val"] / 20
df["rule_8_comp_high"] = df["y_val"] + df["anchor_val"] / 20

def count_matched_in_group(group):
    # 先对当前分组的y_val做排序
    sorted_y = group["y_val"].sort_values().values
    res = []
    for low, high, current_y in zip(group["rule_8_comp_low"], group["rule_8_comp_high"], group["y_val"]):
        # 二分查找定位区间左右边界
        left_idx = sorted_y.searchsorted(low, side="left")
        right_idx = sorted_y.searchsorted(high, side="right")
        # 总数量减去自身1条，得到同组其他符合要求的数量
        res.append((right_idx - left_idx) - 1)
    group["y_val_between"] = res
    return group

# 按anchor分组逐个处理
df = df.groupby("anchor_col", group_keys=False).apply(count_matched_in_group)

额外优化建议 #

• 如果数据集还有更大的扩容空间，可以引入Dask框架做并行分组计算，进一步提升处理效率
• 若anchor_val完全和anchor_col一一对应，可以提前把anchor_val的计算提到分组外，减少重复计算量

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