Pandas分组时为不同列应用不同自定义函数的实现

输入数据 #

annotator  event          interval_presence   duration
3          birds          [0,5]               5
3          birds          [7,9]               10
3          voices         [1,2]               10
3          traffic        [1,7]               7
5          voices         [4,7]               4
5          voices         [5,10]              6
5          traffic        [0,1]               4

其中interval_presence列的元素为Pandas Interval类型。

预期输出 #

annotator  event          interval_presence   duration
3          birds          [[0,5],[7,9]]       7.5
3          voices         [1,2]               10
3          traffic        [1,7]               7
5          voices         [4,10]              5
5          traffic        [0,1]               4

现有代码 #

用户已实现interval_presence列的区间合并逻辑：

data = data.groupby(['annotator', 'event'])['interval_presence'] \
    .apply(pd.arrays.IntervalArray) \
    .apply(piso.union) \
    .reset_index()

解决方案 #

要同时给不同列应用不同处理逻辑，直接用groupby.agg()方法即可，通过字典明确指定每列对应的操作：

import pandas as pd
import pandas_interval as piso

# 把区间合并的逻辑封装成函数
def merge_intervals(intervals):
    return piso.union(pd.arrays.IntervalArray(intervals))

# 执行分组聚合
result = data.groupby(['annotator', 'event']).agg(
    interval_presence=('interval_presence', merge_intervals),
    duration=('duration', 'mean')
).reset_index()

说明 #

• 用字典传递列名和对应处理逻辑：interval_presence列使用封装好的区间合并函数，duration列直接调用内置的mean方法
• 自定义函数merge_intervals把原代码的两步apply逻辑整合，让代码更简洁
• 最后调用reset_index()将分组索引还原为普通列，和预期输出结构完全匹配

内容的提问来源于stack exchange，提问作者M.Tailleur