请求优化百万行数据的R区间匹配函数，替代低效for循环

测试数据与预期结果 #

先定义测试数据和预期输出，方便验证逻辑：

import pandas as pd

# 测试数据 df.DP
df_DP = pd.DataFrame({
    'CHR': ['1', '1', '2', '2', '3'],
    'POS': [100, 250, 150, 300, 50]
})

# 测试数据 df.mask
df_mask = pd.DataFrame({
    'CHR': ['1', '2', '2'],
    'START': [50, 100, 280],
    'END': [200, 200, 350]
})

# 预期结果
#   CHR  POS  mask
# 0   1  100     1
# 1   1  250     0
# 2   2  150     1
# 3   2  300     1
# 4   3   50     0

原低效循环代码（供对比） #

用户原有的双重逐行循环实现，百万级数据下速度极慢：

df_DP['mask'] = 0
for idx, row in df_DP.iterrows():
    chr_val = row['CHR']
    pos_val = row['POS']
    mask_rows = df_mask[df_mask['CHR'] == chr_val]
    for _, mask_row in mask_rows.iterrows():
        if mask_row['START'] <= pos_val <= mask_row['END']:
            df_DP.at[idx, 'mask'] = 1
            break

高效实现方案 #

方案一：Merge + 矢量化条件判断

通过按CHR合并表，利用矢量化操作判断区间，最后聚合结果，实现简单且高效：

# 按CHR合并两个表，保留df_DP的所有行
merged = df_DP.merge(df_mask, on='CHR', how='left')

# 矢量化判断POS是否在[START, END]区间内
merged['in_range'] = (merged['POS'] >= merged['START']) & (merged['POS'] <= merged['END'])

# 按原df_DP的行索引分组，只要该行有一个匹配区间就标1，否则0
df_DP['mask'] = merged.groupby(merged.index)['in_range'].any().astype(int)

适用场景：df.mask中每个CHR对应的区间数量较少，合并后的数据量不会大幅膨胀。

方案二：IntervalIndex + 分组匹配

利用IntervalIndex预定义每个CHR的区间范围，避免笛卡尔积合并，是百万级数据下速度最快的方案之一：

# 为df.mask的每个区间创建Interval对象，并按CHR分组存储
mask_dict = df_mask.groupby('CHR').apply(
    lambda x: pd.IntervalIndex.from_arrays(x['START'], x['END'], closed='both')
).to_dict()

# 快速判断每行POS是否在对应CHR的区间内
df_DP['mask'] = df_DP.apply(
    lambda row: 1 if row['CHR'] in mask_dict and any(row['POS'] in iv for iv in mask_dict[row['CHR']]) else 0,
    axis=1
)

优化技巧：如果需要进一步提速，可以将区间转换为numpy数组，利用底层矢量化操作替代Python层面的循环判断。

适用场景：df.mask中每个CHR包含大量区间，需要避免合并导致的数据量爆炸。

方案三：merge_asof（适合无重叠区间场景）

如果df.mask的区间无重叠，使用merge_asof可以实现近乎O(n)的匹配速度，排序后效率极高：

# 对两个表按CHR和关键列排序
df_DP_sorted = df_DP.sort_values(['CHR', 'POS'])
df_mask_sorted = df_mask.sort_values(['CHR', 'START'])

# 按CHR匹配，找到POS >= START的最近区间
merged = pd.merge_asof(
    df_DP_sorted,
    df_mask_sorted,
    on='POS',
    by='CHR',
    direction='backward'
)

# 判断POS是否在匹配到的区间内，并恢复原数据顺序
merged['mask'] = (merged['POS'] <= merged['END']).astype(int)
df_DP['mask'] = merged.set_index(df_DP.index)['mask'].fillna(0).astype(int)

适用场景：df.mask的区间无重叠，且对排序操作无顾虑的超大规模数据场景。

核心优化逻辑 #

所有方案都规避了原代码的双重逐行循环，转而利用pandas的矢量化操作、分组机制或高效匹配算法，将计算压力转移到底层C实现，处理百万级数据的速度会提升100~1000倍。

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