如何根据比较两列中的值来组合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的行产生了重叠，那么它们就会被组合起来。(在这里，指数是不重要的。)有没有一种聪明或有效的方法可以在不进行大量复杂的迭代的情况下做到这一点？(我的实际数据可能有多达数百万行。)

上面例子的最终结果应该是。

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

Answer 1

一种可能的方法是通过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的解答。

Answer 2

一种方法是创建一个函数来合并区间，然后 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