Как я могу объединить строки в кадре данных pandas на основе сравнения значений в двух столбцах?

Question

Рассмотрим pandas фрейм данных, например:

df = pd.DataFrame({'id': ['001', '001', '002', '002', '003', '003', '004', '004', '005', '005'], 
                      'start': [1, 200, 200, 1, 1, 200, 200, 1, 1, 1000],
                      'end': [1000, 500, 500, 1000, 500, 1000, 1000, 500, 500, 2000]})
or 
    id  start   end
0  001      1  1000
1  001    200   500
2  002    200   500
3  002      1  1000
4  003      1   500
5  003    200  1000
6  004    200  1000
7  004      1   500
8  005      1   500
9  005   1000  2000

Я хотел бы получить pandas фрейм данных, так что если начало и конец производят перекрытие для строки с данным идентификатором, тогда они совмещены. (Индексы здесь не важны.) Есть ли умный или эффективный способ сделать это, не прибегая к множеству сложных итераций? (Мои фактические данные могут содержать до миллионов строк.)

Конечный результат из приведенного выше примера должен быть:

   id  start   end
  001      1  1000
  002      1  1000
  003      1  1000
  004      1  1000
  005      1   500
  005   1000  2000

Answer 1

Один из способов - создать функцию для объединения интервалов, а затем groupby и apply:

def merge(l):
    l = sorted(l, key=lambda x: x[0])
    merged = []
    for i in l:
        if not merged or merged[-1][1] < i[0]:
            merged.append(i)
        else:
            merged[-1][1] = max(merged[-1][1], i[1])
    return merged

print (df.groupby("id").apply(lambda d: merge(d[["start","end"]].values)).explode())

id
001       [1, 1000]
002       [1, 1000]
003       [1, 1000]
004       [1, 1000]
005        [1, 500]
005    [1000, 2000]
dtype: object

Answer 2

Одним из возможных способов сделать это может быть группировка по столбцу id с последующим применением функции для интервалов слияния:

import pandas as pd

# Load or create the dataframe df.

def merge_intervals(group):
    l = zip(group['start'], group['end'])
    merged = []
    for i in sorted(l):
        if not merged or merged[-1][1] < i[0]:
            merged.append(list(i))
        else:
            merged[-1][1] = max(merged[-1][1], i[1])
    start, end = zip(*[(x[0], x[1]) for x in merged])
    return pd.DataFrame({
        'id':    group['id'][0],
        'start': start,
        'end':   end
    })

df_new = df.groupby(df['id'], as_index=False).apply(merge_intervals)

В вашем примере результат выглядит следующим образом

      id  start   end
0 0  001      1  1000
1 0  002      1  1000
2 0  003      1  1000
3 0  004      1  1000
4 0  005      1   500
  1  005   1000  2000

Создается MultiIndex DataFrame, в котором будут те же столбцы.

Спасибо @ scott-boston за обнаружение ошибки и @ henry-yik за ответ.