如何在Pandas中按多条件分组合并DataFrame行？

需求说明 #

• 按user字段分组
• 按start_at_min字段排序
• 合并条件分三个场景：
  • 场景A：当后一行start_at_min≤200时，若后一行start_at_min - 前一行stop_at_min <5则合并
  • 场景B：当200<后一行start_at_min<400时，合并阈值改为3
  • 场景C：当后一行start_at_min>400时，绝不合并

示例DataFrame #

user  start_at_min  stop_at_min
0     1           100          150  
1     1           152          201 # row0与row1合并
2     1           205          260 # row1与row2不合并 -> start_at_min超过200，阈值为3
3     2            65          100 # 不合并
4     2           200          265 # 不合并
5     2           300          451 # 不合并
6     2           452          460 # 不合并 -> start_at_min超过400，绝不合并

预期输出 #

user  start_at_min  stop_at_min
0     1           100          201 # row0与row1合并后的结果
2     1           205          260 # row1与row2不合并
3     2            65          100 # 无变化
4     2           200          265 # 无变化
5     2           300          451 # 无变化
6     2           452          460 # 无变化

现有代码及问题 #

我写了一个combine_rows函数尝试实现两行合并逻辑，但无法正确应用到分组排序后的DataFrame上：

合并函数 #

def combine_rows (s1:pd.Series, s2:pd.Series):
  # 若后一行start_at_min与前一行stop_at_min差值小于5则合并
  if s2['start_at_min'] - s1['stop_at_min'] <5: 
     return pd.Series({
         'user': s1['user'],
         'start_at_min': s1['start_at_min'],
         'stop_at_min' : s2['stop_at_min']
         })
  else: 
    return pd.concat([s1,s2],axis=1).T

尝试调用代码（无效） #

df.groupby('user').sort_values(by=['start_at_min']).apply(combine_rows) # 无法正常工作

完整测试代码 #

import pandas as pd 
import numpy as np


df = pd.DataFrame({
    "user"       :  [1, 1, 2,2],
    'start_at_min': [60, 101, 65, 200], 
    'stop_at_min' : [100, 135, 100, 265] 
})

def combine_rows (s1:pd.Series, s2:pd.Series):
  # 若后一行start_at_min与前一行stop_at_min差值小于5则合并
  if s2['start_at_min'] - s1['stop_at_min'] <5: 
     return pd.Series({
         'user': s1['user'],
         'start_at_min': s1['start_at_min'],
         'stop_at_min' : s2['stop_at_min']
         })
  else: 
    return pd.concat([s1,s2],axis=1).T

df.groupby('user').sort_values(by=['start_at_min']).apply(combine_rows) # 无法正常工作

解决方案 #

问题核心：分组后apply默认传入的是整个组的DataFrame，而非两两Series。需要遍历分组内的行，按逻辑逐步合并。

正确实现函数 #

def merge_group(group):
    # 先按start_at_min排序
    group = group.sort_values('start_at_min').reset_index(drop=True)
    merged = []
    # 初始化当前合并行
    current_row = group.iloc[0].copy()
    
    for idx in range(1, len(group)):
        next_row = group.iloc[idx]
        # 根据next_row的start_at_min确定阈值
        if next_row['start_at_min'] <= 200:
            threshold = 5
        elif 200 < next_row['start_at_min'] < 400:
            threshold = 3
        else:
            # 超过400，直接存入结果，更新当前行
            merged.append(current_row)
            current_row = next_row.copy()
            continue
        
        # 检查是否满足合并条件
        if next_row['start_at_min'] - current_row['stop_at_min'] < threshold:
            # 合并：更新当前行的stop_at_min
            current_row['stop_at_min'] = next_row['stop_at_min']
        else:
            # 不合并，存入结果，更新当前行
            merged.append(current_row)
            current_row = next_row.copy()
    # 加入最后一行
    merged.append(current_row)
    return pd.DataFrame(merged)

调用方式 #

result = df.groupby('user', group_keys=False).apply(merge_group)
print(result)

测试结果 #

针对提供的测试DataFrame，输出如下：

user  start_at_min  stop_at_min
0     1            60          135  # 60-100与101-135合并（差值1<5）
1     2            65          100
2     2           200          265

完全符合预期的合并逻辑。

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