Использование 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. df2 подмножество df1
   • 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]

Я пытаюсь использовать 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'
)

находит ближайший df1.end после 800 в df2 [650 800], что составляет 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.end, только если df1.start в этом диапазоне больше 800 (950 в этом дело)"? С таким уровнем сложности, может быть, есть другая функция, лучше подходящая для этой задачи?

Примечания:

• Диапазоны в df1 могут перекрывать друг друга, но никогда не бывают идентичными.
• Диапазоны в df2 могут перекрывать друг друга, но никогда не бывают идентичными.
• df1 и df2 имеют более 200 000 строк каждая.
• 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. Чтобы найти суперсет, его можно запустить в обратном порядке. Для простоты он принимает список диапазонов вместо DataFrame s. Преобразование является простым.

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