如何在Pandas中无需循环高效比对两个DataFrame并筛选特定行对

问题需求 #

比对just_f和just_r两个DataFrame，筛选满足以下条件的行对：

• sample、chr、family字段值完全匹配
• just_r的pos值处于just_f的pos到pos+1000区间内

当前采用itertuples双重循环实现，数千行数据时耗时过长，需优化。

输入数据示例 #

just_f

sample  chr pos strand  family  order   support comment frequency
2   NC_025812.2 9831    .   Tourist|7   Tourist F   -   0,562
2   NC_025812.2 12038   .   Tourist|7   Tourist F   -   1,000
5   NC_025812.2 12040   .   Tourist|7   Tourist F   -   1,000
12  NC_025812.2 12042   .   Tourist|7   Tourist F   -   1,000
11  NC_025812.2 30758   .   uc|32   uc  F   -   0,547
12  NC_025812.2 49544   .   uc|10   uc  F   -   0,112
11  NC_025812.2 56184   .   hAT|9   hAT F   -   0,997
5   NC_025812.2 56246   .   hAT|9   hAT F   -   0,756
3   NC_025812.2 56265   .   hAT|9   hAT F   -   1,000
12  NC_025812.2 56268   .   hAT|9   hAT F   -   1,000

just_r

5   NC_025812.2 12396   .   Tourist|7   Tourist R   -   0,975
2   NC_025812.2 12433   .   Tourist|7   Tourist R   -   0,935
12  NC_025812.2 12478   .   Tourist|7   Tourist R   -   0,887
12  NC_025812.2 28943   .   Tourist|7   Tourist R   -   0,610
5   NC_025812.2 28947   .   Tourist|7   Tourist R   -   0,490
2   NC_025812.2 51483   .   Mutator|24  Mutator R   -   0,422
5   NC_025812.2 56713   .   hAT|9   hAT R   -   0,925
11  NC_025812.2 56737   .   hAT|9   hAT R   -   1,000
3   NC_025812.2 56778   .   hAT|9   hAT R   -   0,891
12  NC_025812.2 56800   .   hAT|9   hAT R   -   0,965

预期输出（f_r_pairs） #

sample  chr pos strand  family  order   support comment frequency
2   NC_025812.2 12038   .   Tourist|7   Tourist F   -   1.0
2   NC_025812.2 12433   .   Tourist|7   Tourist R   -   0.935
5   NC_025812.2 12040   .   Tourist|7   Tourist F   -   1.0
5   NC_025812.2 12396   .   Tourist|7   Tourist R   -   0.975
12  NC_025812.2 12042   .   Tourist|7   Tourist F   -   1.0
12  NC_025812.2 12478   .   Tourist|7   Tourist R   -   0.887
11  NC_025812.2 56184   .   hAT|9   hAT F   -   0.997
11  NC_025812.2 56737   .   hAT|9   hAT R   -   1.0
5   NC_025812.2 56246   .   hAT|9   hAT F   -   0.756
5   NC_025812.2 56713   .   hAT|9   hAT R   -   0.925
3   NC_025812.2 56265   .   hAT|9   hAT F   -   1.0
3   NC_025812.2 56778   .   hAT|9   hAT R   -   0.891
12  NC_025812.2 56268   .   hAT|9   hAT F   -   1.0
12  NC_025812.2 56800   .   hAT|9   hAT R   -   0.965

原低效代码 #

import pandas as pd

df_raw = pd.read_csv('1-DH-to-12-RO.NC_teinsertions.txt', sep="\t", decimal=',')
df_sort = df_raw.sort_values(by=['chr', 'pos', 'sample'])

just_f = df_sort[(df_sort["support"] == 'F')]
just_r = df_sort[(df_sort["support"] == 'R')]

f_r_pairs = pd.DataFrame(columns=just_f.columns)

# choosing rows for reference TE insertions (having pairs with F and R in range 1000 bp)
for f in just_f.itertuples():
    for r in just_r.itertuples():
        if f.sample == r.sample and f.chr == r.chr and f.family == r.family and r.pos in range(f.pos, f.pos + 1000):
            f_r_pairs = f_r_pairs.append(pd.DataFrame([f]))
            f_r_pairs = f_r_pairs.append(pd.DataFrame([r]))

优化方案 #

方案一：合并+条件过滤（简洁高效） #

利用Pandas内置的merge操作（C级实现）完成分组匹配，再过滤区间条件，避免Python循环。

import pandas as pd

# 读取并预处理数据
df_raw = pd.read_csv('1-DH-to-12-RO.NC_teinsertions.txt', sep="\t", decimal=',')
df_sort = df_raw.sort_values(by=['chr', 'pos', 'sample'])

just_f = df_sort[df_sort["support"] == 'F'].copy()
just_r = df_sort[df_sort["support"] == 'R'].copy()

# 重命名pos字段，避免合并后冲突
just_f.rename(columns={'pos': 'f_pos'}, inplace=True)
just_r.rename(columns={'pos': 'r_pos'}, inplace=True)

# 按共同字段合并，得到所有候选对
merged = pd.merge(just_f, just_r, on=['sample', 'chr', 'family'], suffixes=('_f', '_r'))

# 过滤pos区间条件
filtered = merged[(merged['r_pos'] >= merged['f_pos']) & (merged['r_pos'] <= merged['f_pos'] + 1000)]

# 整理成预期的行对格式
# 提取F行数据
f_rows = filtered[[col for col in filtered.columns if not col.endswith('_r')]]
f_rows.rename(columns={'f_pos': 'pos'}, inplace=True)
# 提取R行数据
r_cols = [col.replace('_r', '') for col in filtered.columns if col.endswith('_r') or col in ['sample', 'chr', 'family']]
r_rows = filtered[r_cols]
r_rows.rename(columns={'r_pos': 'pos'}, inplace=True)

# 合并并排序，与预期输出一致
f_r_pairs = pd.concat([f_rows, r_rows]).sort_values(by=['sample', 'chr', 'pos']).reset_index(drop=True)

方案二：分组后区间匹配（大数据量更优） #

先按sample、chr、family分组，仅在同组内做区间匹配，减少不必要的比对，适合数据量较大的场景。

import pandas as pd

# 读取并预处理数据
df_raw = pd.read_csv('1-DH-to-12-RO.NC_teinsertions.txt', sep="\t", decimal=',')
df_sort = df_raw.sort_values(by=['chr', 'pos', 'sample'])

just_f = df_sort[df_sort["support"] == 'F']
just_r = df_sort[df_sort["support"] == 'R']

# 预先对just_r按分组字段分组，提升查询效率
r_groups = just_r.groupby(['sample', 'chr', 'family'])

pairs_list = []

# 遍历just_f的每个分组
for (sample, chr_val, family), f_group in just_f.groupby(['sample', 'chr', 'family']):
    # 若该分组在just_r中不存在，直接跳过
    if (sample, chr_val, family) not in r_groups.groups:
        continue
    r_group = r_groups.get_group((sample, chr_val, family))
    
    # 对组内每个F行，匹配符合区间条件的R行
    for _, f_row in f_group.iterrows():
        # 矢量化筛选，比循环快
        mask = (r_group['pos'] >= f_row['pos']) & (r_group['pos'] <= f_row['pos'] + 1000)
        matched_r = r_group[mask]
        if not matched_r.empty:
            pairs_list.append(f_row)
            # 将匹配的R行转为字典列表加入结果
            pairs_list.extend(matched_r.to_dict('records'))

# 转换为DataFrame并排序
f_r_pairs = pd.DataFrame(pairs_list).sort_values(by=['sample', 'chr', 'pos']).reset_index(drop=True)

优化原理 #

• 原方案的双重Python循环是O(n²)复杂度，数千行数据会产生百万级比对操作，效率极低
• 优化方案利用Pandas的矢量化运算和内置分组/合并操作，底层基于C实现，运算速度提升几个数量级
• 方案一适合数据分布均匀的场景，代码简洁；方案二通过分组减少无效比对，适合大数据量或分组稀疏的场景

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