计算时长均值后，如何保留Pandas DataFrame中的cases列？

生成DataFrame的代码 #

import pandas as pd
from io import StringIO

df = """
 b_id          duration1                  duration2                          user  cases
 
 366           NaN                        38 days 22:05:06.807430            Test  4
 367           0 days 00:00:05.285239     NaN                                Test  4
 368           NaN                        NaN                                Test  4
 366           NaN                        38 days 22:05:06.807430            Test  4
 
 466           NaN                        38 days 22:05:06.807430            Tom   3
 467           0 days 00:00:05.285239     NaN                                Tom   3
 467           0 days 00:00:05.285239     NaN                                Tom   3
"""
df= pd.read_csv(StringIO(df.strip()), sep='\s\s+', engine='python')

原计算代码及问题 #

原代码仅计算duration列的均值，导致cases列丢失：

out = (df
   .set_index('user')
   .filter(like='duration')
   .apply(pd.to_timedelta)
   .groupby(level=0).mean()
   .reset_index()
 )

输出结果：

user    duration1   duration2
0   Test    0 days 00:00:05.285239  38 days 22:05:06.807430
1   Tom     0 days 00:00:05.285239  38 days 22:05:06.807430

期望输出 #

希望保留cases列，输出如下：

user    duration1   duration2                               cases
0   Test    0 days 00:00:05.285239  38 days 22:05:06.807430     4
1   Tom     0 days 00:00:05.285239  38 days 22:05:06.807430     3

解决方案 #

由于每个user对应的cases值唯一，可通过以下两种方式实现：

方法一：分组时同时聚合cases列

直接在分组聚合阶段，对duration列取均值，cases列取任意单一值（如first，因同一用户cases值一致）：

out = (df
       # 先将duration列转为时间差类型
       .assign(
           duration1=pd.to_timedelta(df['duration1']),
           duration2=pd.to_timedelta(df['duration2'])
       )
       # 按user分组，指定各列的聚合方式
       .groupby('user', as_index=False)
       .agg(
           duration1=('duration1', 'mean'),
           duration2=('duration2', 'mean'),
           cases=('cases', 'first')
       )
)

方法二：先计算均值再合并cases映射

先计算duration的均值，再通过用户与cases的唯一映射表合并结果：

# 计算duration列的均值
duration_mean = (df
                 .set_index('user')
                 .filter(like='duration')
                 .apply(pd.to_timedelta)
                 .groupby(level=0).mean()
                 .reset_index()
                )
# 获取user与cases的唯一对应关系
user_cases_map = df[['user', 'cases']].drop_duplicates()
# 合并两个结果表
out = pd.merge(duration_mean, user_cases_map, on='user')

两种方法都能得到期望的输出，可根据个人习惯选择。

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