大体积Pandas DataFrame识别同染色体重叠基因组区域高效方案

问题背景 #

现有存储基因组区域信息的Pandas DataFrame，包含chromosome、start、stop、hg_38_locs（区域标识）、main_category（标签）字段，数据量达300万行。需完成：

• 识别同一染色体上的重叠区域对
• 汇总对应区域的标签
• 方案需高效，禁止使用for循环遍历

示例数据 #

import pandas as pd

# 示例DataFrame
data = {
    'chromosome': ['chr1', 'chr1', 'chr1', 'chr1', 'chr1'],
    'start': [10, 15, 35, 45, 55],
    'stop': [20, 25, 55, 56, 60],
    'hg_38_locs': ['chr1:10-20', 'chr1:15-25', 'chr1:35-55', 'chr1:45-56', 'chr1:55-60'],
    'main_category': ['label1', 'label2', 'label2', 'label3', 'label1']
}
df = pd.DataFrame(data)

预期输出（修正非重叠错误后） #

注：原示例输出包含非重叠区域对，此处按真实重叠逻辑修正；若需保留原输出逻辑，可调整过滤条件。

高效实现方案 #

方案1：Pandas自连接+条件过滤（无需额外库） #

利用Pandas向量化运算替代循环，通过分组自连接减少无效计算：

def find_overlapping_pairs(df):
    result_dfs = []
    # 按染色体分组处理，避免跨染色体无效配对
    for _, group in df.groupby('chromosome'):
        # 自连接生成所有区域对，排除自身配对
        pair_df = pd.merge(group, group, on='chromosome', suffixes=('_x', '_y'))
        pair_df = pair_df[pair_df['hg_38_locs_x'] != pair_df['hg_38_locs_y']]
        
        # 过滤重叠区域：A的起始 < B的结束 且 A的结束 > B的起始
        overlap_mask = (pair_df['start_x'] < pair_df['stop_y']) & (pair_df['stop_x'] > pair_df['start_y'])
        pair_df = pair_df[overlap_mask]
        
        # 避免重复记录（只保留A-B，不重复记录B-A）
        pair_df = pair_df[pair_df['hg_38_locs_x'] < pair_df['hg_38_locs_y']]
        
        # 构造输出格式
        pair_df['overlapping_regions'] = pair_df.apply(lambda row: (row['hg_38_locs_x'], row['hg_38_locs_y']), axis=1)
        pair_df['overlapping_labels'] = pair_df.apply(lambda row: (row['main_category_x'], row['main_category_y']), axis=1)
        
        result_dfs.append(pair_df[['overlapping_regions', 'overlapping_labels']])
    
    return pd.concat(result_dfs).reset_index(drop=True)

# 运行
final_df = find_overlapping_pairs(df)
print(final_df)

方案2：IntervalTree（适配300万行超大规模数据） #

采用区间树结构实现O(n log n)级别的高效查询，需先安装依赖：

pip install intervaltree

实现代码：

from intervaltree import IntervalTree

def find_overlapping_pairs_intervaltree(df):
    result = []
    seen_pairs = set()  # 去重，避免重复记录同一配对
    
    for chrom, group in df.groupby('chromosome'):
        # 构建区间树，存储区间与对应区域标识、标签
        tree = IntervalTree()
        for _, row in group.iterrows():
            tree.addi(row['start'], row['stop'], (row['hg_38_locs'], row['main_category']))
        
        # 遍历每个区域，查询重叠区间
        for _, row in group.iterrows():
            current_loc = row['hg_38_locs']
            current_label = row['main_category']
            # 查询当前区域的所有重叠区间
            overlaps = tree.search(row['start'], row['stop'])
            
            for interval in overlaps:
                overlap_loc, overlap_label = interval.data
                # 排除自身，且确保配对只记录一次
                pair_key = tuple(sorted((current_loc, overlap_loc)))
                if current_loc != overlap_loc and pair_key not in seen_pairs:
                    seen_pairs.add(pair_key)
                    result.append({
                        'overlapping_regions': (current_loc, overlap_loc),
                        'overlapping_labels': (current_label, overlap_label)
                    })
    
    return pd.DataFrame(result)

# 运行
final_df = find_overlapping_pairs_intervaltree(df)
print(final_df)

方案选型建议 #

• 方案1：无需额外依赖，适合百万级以内数据，代码简洁易维护；
• 方案2：查询效率更优，针对300万行超大规模数据，能显著降低运行时间。

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