如何使Pandas GroupBy Transform结果填充DataFrame所有行？

我在使用Pandas做GroupBy转换时，当前代码仅给被mask标记的行赋值了转换后的值，但我希望value_transform列的所有单元格都填充对应日期的计算值。原代码及预期输出如下：

import pandas as pd

data = {
    'dates': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03'],
    'flag': [True, True, False, False, False, True, True, False, True],
    'value': [10, 10, 20, 30, 40, 50, 50, 50, 60],
}

df = pd.DataFrame(data)
mask = df['flag'] == True
df['value_transform'] = df[mask].groupby('dates')['value'].transform(lambda x: max(x) - min(x))

预期输出：

dates   flag  value  value_transform
0  2023-01-01   True     10    40.0
1  2023-01-02   True     10     0.0
2  2023-01-03  False     20    10.0
3  2023-01-01  False     30    40.0
4  2023-01-02  False     40     0.0
5  2023-01-03   True     50    10.0
6  2023-01-01   True     50    40.0
7  2023-01-02  False     50     0.0
8  2023-01-03   True     60    10.0

解决方案 #

方法一：先计算日期统计映射，再批量赋值 #

先筛选出mask标记的行，按日期分组计算max(value) - min(value)，得到日期到计算值的映射关系，再通过map方法给所有行的value_transform列赋值：

import pandas as pd

data = {
    'dates': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03'],
    'flag': [True, True, False, False, False, True, True, False, True],
    'value': [10, 10, 20, 30, 40, 50, 50, 50, 60],
}

df = pd.DataFrame(data)
mask = df['flag'] == True

# 计算每个日期对应的max-min值
date_calc = df[mask].groupby('dates')['value'].agg(lambda x: x.max() - x.min())
# 给所有行赋值对应日期的计算值
df['value_transform'] = df['dates'].map(date_calc)

print(df)

方法二：在groupby的transform中直接针对mask行计算 #

通过groupby的transform方法，在每个分组内仅筛选mask标记的行进行计算，结果会自动填充到所有行：

import pandas as pd

data = {
    'dates': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03', '2023-01-01', '2023-01-02', '2023-01-03'],
    'flag': [True, True, False, False, False, True, True, False, True],
    'value': [10, 10, 20, 30, 40, 50, 50, 50, 60],
}

df = pd.DataFrame(data)
mask = df['flag'] == True

df['value_transform'] = df.groupby('dates')['value'].transform(
    lambda x: x[mask.loc[x.index]].max() - x[mask.loc[x.index]].min()
)

print(df)

两种方法都能得到预期输出，方法一的计算效率更高，适合处理大规模数据集。

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