Python中分组DataFrame后按条件为每行计算自定义统计值的方法

问题背景 #

现有如下Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(
         {"CALDT": ["1980-01-31", "1980-02-28", "1980-03-31",
                    "1980-01-31", "1980-02-28", "1980-03-31",
                    "1980-01-31"],
          "ID": [1, 1, 1, 
                 2, 2, 2,
                 3],
          "Return": [0.02, 0.05, 0.10,
                     0.05, -0.02, 0.03,
                     -0.03]
          })

df['Year'] = pd.to_datetime(df['CALDT']).dt.year

需求：按ID和Year分组，若分组内的存活月份（即唯一CALDT的数量）≥2，则基于Return列计算均值和中位数并乘以12，将结果赋值给该分组的每一行；否则赋值为NaN。

预期输出如下：

df_new = pd.DataFrame(
    {"CALDT": ["1980-01-31", "1980-02-28", "1980-03-31",
               "1980-01-31", "1980-02-28", "1980-03-31",
               "1980-01-31"],
     "Year": [1980, 1980, 1980,
              1980, 1980, 1980,
              1980],
     "ID": [1, 1, 1, 
            2, 2, 2,
            3],
     "Return": [0.02, 0.05, 0.10,
                0.05, -0.02, 0.03,
                -0.03],
     "Mean_Return": [0.68, 0.68, 0.68,
                     0.24, 0.24, 0.24,
                     np.nan],
     "Median_Return": [0.60, 0.60, 0.60,
                       0.36, 0.36, 0.36,
                       np.nan]
     })

在R中可以通过tidyverse的group_by轻松实现：

df = df %>% 
  mutate(Year = year(CALDT)) %>% 
  group_by(CRSP_FUNDNO, Year) %>% 
  mutate(months_alive = length(unique(CALDT))) %>% 
  mutate(mean = case_when(months_alive >= 2 ~ mean(Return)*12,
                          .default = NA)) %>% 
  mutate(median = case_when(months_alive >= 2 ~ median(Return)*12,
                            .default = NA))

Python实现方案 #

以下是对应Pandas的实现代码，逻辑与R版本完全对齐：

步骤1：计算每个分组的存活月份数 #

通过groupby.transform将分组内唯一CALDT的数量广播到每一行：

# 计算每个ID-年份分组的存活月份数
df['months_alive'] = df.groupby(['ID', 'Year'])['CALDT'].transform(lambda x: x.nunique())

步骤2：计算条件化的均值与中位数 #

同样使用transform计算分组内的均值和中位数并乘以12，再通过where方法筛选满足条件的结果：

# 计算年度化均值，仅保留存活月份≥2的分组结果
df['Mean_Return'] = df.groupby(['ID', 'Year'])['Return'].transform(lambda x: x.mean() * 12).where(df['months_alive'] >= 2)

# 计算年度化中位数，仅保留存活月份≥2的分组结果
df['Median_Return'] = df.groupby(['ID', 'Year'])['Return'].transform(lambda x: x.median() * 12).where(df['months_alive'] >= 2)

步骤3：整理结果（可选） #

如果不需要中间列months_alive，可以直接移除：

# 移除中间列，得到最终结果
df_new = df.drop('months_alive', axis=1)

执行上述代码后，df_new的结果将与预期输出完全一致。

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