使用 pandas merge_asof() 识别范围关系

Question

给定以下两个代表范围的数据框：

df1 =

  start   end
0   200   300
1   600   900
2   950  1050

df2 =

  start   end
0   350   550
1   650   800
2   900  1100

它们可以这样表示：

df1  [200 300]            [600    900] [950 1050]
df2            [350  550]   [650 800] [900   1100]

我的任务是确定df1 和df2 范围之间的四种不同类型的关系：

1. df2df1 的子集
   • df2 [650 800] df1 [600 900] 的子集
2. df2 的超集 df1
   • df2 [900 1100] 的超集 df1 [950 1050]
3. df2 在df1 之后（最近邻，不包括子集/超集）
   • df2 [350 550] 在df1 [200 300] 之后
   • df2 [900 1100] 在df1 [600 900] 之后
4. df2 在df1 之前（最近邻，不包括子集/超集）
   • df2 [350 550] 在 df1 [600 900] 之前
   • df2 [650 800] 在df1 [950 1050] 之前

我正在尝试使用从this answer 学到的merge_asof()，但由于超集/子集关系添加的复杂性，它不起作用，例如：

# Create "before" condition
df_before = pd.merge_asof(
    df2.rename(columns={col:f'before_{col}' for col in df2.columns}).sort_values('before_end'),
    df1.assign(before_end=lambda x: x['end']).sort_values('before_end'),
    on='before_end',
    direction='forward'
).query('end > before_end')

print(df_before)

输出：

  before_start  before_end  start    end
0          350         550  600.0  900.0
1          650         800  600.0  900.0

目标输出：

  before_start  before_end  start     end
0          350         550  600.0   900.0
1          650         800  950.0  1050.0

问题是

pd.merge_asof(
    df2.rename(columns={col:f'before_{col}' for col in df2.columns}).sort_values('before_end'),
    df1.assign(before_end=lambda x: x['end']).sort_values('before_end'),
    on='before_end',
    direction='forward'
)

在df2 [650 800]中找到800之后最接近的df1.end，即df1 [600 900]：

  before_start  before_end  start    end
0          350         550  600.0  900.0
1          650         800  600.0  900.0
2          900        1100    NaN    NaN

是否可以根据特定条件执行merge_asof() 来查找最近的值，例如“仅当该范围内的df1.start 大于800（本例中为950）时才查找最近的df1.end” ?有了这种复杂程度，也许还有其他功能更适合这项任务？

注意事项：

• df1 中的范围可以相互重叠，但绝不相同。
• df2 中的范围可以相互重叠，但绝不相同。
• df1 和 df2 各有超过 20 万行。
• df1 和 df2 的行数不同。
• 关系与df1 相关，因此df1 中的每一行只需要一个匹配项，每行最多有四个关系。鉴于上面提供的数据，合并回df1 后的最终输出将如下所示：

df1 =

  start   end  subset_start  subset_end  superset_start  superset_end  before_start  before_end  after_start  after_end
0   200   300           NaN         NaN             NaN           NaN           NaN         NaN        350.0      550.0
1   600   900         650.0       800.0             NaN           NaN         350.0       550.0        900.0     1100.0
2   950  1050           NaN         NaN           900.0        1100.0         650.0       800.0          NaN        NaN

Answer 1

可以使用pd.merge_asof查找前后选项。

before_df = pd.merge_asof(df1, df2, left_on='start', right_on='end', suffixes=['', '_before'])
before_df
#    start   end  start_before  end_before
# 0    200   300           NaN         NaN
# 1    600   900         350.0       550.0
# 2    950  1050         650.0       800.0

after_df = pd.merge_asof(df2, df1, left_on='start', right_on='end', suffixes=['_after', ''])
#    start_after  end_after  start  end
# 0          350        550    200  300
# 1          650        800    200  300
# 2          900       1100    600  900

但要使其工作或子集和超集计算并不容易。对于那些，我会使用这种可以一次性工作的算法。

def range_intersect(lh_ranges, rh_ranges): 
    all_ranges = sorted(
        [(b, e, 'lh') for b, e in lh_ranges] +
        [(b, e, 'rh') for b, e in rh_ranges]
    ) 

    res = [] 
    max_b, max_e = None, None 
    for b, e, which in all_ranges: 
        if which == 'rh': 
            if max_e is None or e > max_e: 
                max_b, max_e = b, e 
        elif max_e is not None and e <= max_e: 
            res.append((b, e, max_b, max_e)) 

    return res

这会找到lh 的元素，它们是rh 中元素的子集。要查找超集，可以反向运行。为简单起见，它采用范围列表而不是DataFrames。转换很简单。

lh = df1.to_dict('split')['data']
rh = df2.to_dict('split')['data']

lh
# [[200, 300], [600, 900], [950, 1050]]

rh                                                                                                                                                                                                                                  
# [[350, 550], [650, 800], [900, 1100]]

之后，您想要的结果 DataFrame 只需几个合并即可。

# Compute supersets, then run in reverse to get the subsets.
superset_df = pd.DataFrame(range_intersect(lh, rh), columns=['start', 'end', 'start_superset', 'end_superset'])
subset_df = pd.DataFrame(range_intersect(rh, lh), columns=['start_subset', 'end_subset', 'start', 'end'])

# Merge all the results together.
result = df1.merge(subset_df, how='left').merge(superset_df, how='left').merge(before_df, how='left').merge(after_df, how='left')

# The reversed operations, after and subset, can have many matches in df1.
result.drop_duplicates(['start', 'end'])
#    start   end  start_subset  end_subset  start_superset  end_superset  start_before  end_before  start_after  end_after
# 0    200   300           NaN         NaN             NaN           NaN           NaN         NaN        350.0      550.0
# 2    600   900         650.0       800.0             NaN           NaN         350.0       550.0        900.0     1100.0
# 3    950  1050           NaN         NaN           900.0        1100.0         650.0       800.0          NaN        NaN