如何移除dataframe中同id下字符索引区间重叠的较短实体行？

核心思路 #

同id分组后，优先保留更长的实体区间，再依次判断后续短区间是否与已保留的区间重叠，重叠则丢弃，不重叠则保留。该思路时间复杂度为O(n log n)，性能远高于暴力匹配方案。

可运行代码示例 #

首先导入依赖并构造测试数据：

import pandas as pd

# 构造示例dataframe
data = [
    [1, "very sunny day", 11, 26, [4,5,6]],
    [1, "shining", 37, 45, [10]],
    [1, "sunny", 16, 21, [5]],
    [2, "the red balloon", 9, 25, [3,4,5]],
    [2, "playing", 29, 37, [7]],
    [2, "red", 13, 16, [4]]
]
df = pd.DataFrame(data, columns=["id", "data", "start_char_index", "end_char_index", "token_position"])

基础版（10万行以内适用） #

def filter_overlapping_entities(group):
    # 先按区间长度降序排序，长度相同时按起始位置升序，保证长区间优先被保留
    group = group.sort_values(
        by=["length", "start_char_index"], 
        ascending=[False, True]
    ).reset_index(drop=True)
    keep_rows = []
    kept_intervals = []
    for _, row in group.iterrows():
        s, e = row["start_char_index"], row["end_char_index"]
        # 判断是否和已保留区间重叠：当前区间起始 < 已存区间结束 且 当前区间结束 > 已存区间起始
        is_overlap = any(ks < e and ke > s for ks, ke in kept_intervals)
        if not is_overlap:
            keep_rows.append(row)
            kept_intervals.append((s, e))
    return pd.DataFrame(keep_rows)

# 计算每个实体的区间长度
df["length"] = df["end_char_index"] - df["start_char_index"]
# 按id分组应用过滤规则
result = df.groupby("id", group_keys=False).apply(filter_overlapping_entities).drop(columns="length").reset_index(drop=True)

高性能版（百万级数据适用，基于IntervalTree优化重叠判断） #

先安装依赖：

pip install intervaltree

代码如下：

from intervaltree import IntervalTree

def filter_overlapping_entities_fast(group):
    group = group.sort_values(
        by=["length", "start_char_index"], 
        ascending=[False, True]
    ).reset_index(drop=True)
    keep_rows = []
    interval_tree = IntervalTree()
    for _, row in group.iterrows():
        s, e = row["start_char_index"], row["end_char_index"]
        # 直接调用IntervalTree的重叠判断接口，时间复杂度为O(log k)，k为已保留的区间数
        if not interval_tree.overlaps(s, e):
            keep_rows.append(row)
            interval_tree.addi(s, e)
    return pd.DataFrame(keep_rows)

df["length"] = df["end_char_index"] - df["start_char_index"]
result = df.groupby("id", group_keys=False).apply(filter_overlapping_entities_fast).drop(columns="length").reset_index(drop=True)

输出结果验证 #

运行上述代码后得到的result即为期望的输出：

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