本文旨在解决 Pandas DataFrame 分组后，根据组内数据计算特定统计量（如均值和中位数），并将结果应用回原始DataFrame的每行数据的问题。我们将演示如何使用 groupby()、transform() 和 where() 函数，结合条件判断，高效地实现这一目标，并生成新的包含计算结果的列。
在数据分析中，经常需要对DataFrame进行分组，并对每个组应用特定的函数。一个常见的需求是，根据分组后的数据计算统计量，并将这些统计量作为新的列添加回原始DataFrame。例如，我们可能需要计算每个用户的平均消费金额，并将该平均值添加到每个用户的消费记录中。如果还需要根据组的大小应用不同的计算逻辑，情况会变得更加复杂。本文将介绍如何使用 Pandas 的 groupby()、transform() 和 where() 函数来解决这类问题。
准备数据
首先，我们创建一个示例DataFrame，模拟包含日期、ID和收益率的数据：
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['CALDT'] = pd.to_datetime(df['CALDT'])
df['Year'] = df['CALDT'].dt.year
print(df)这段代码创建了一个包含日期 (CALDT)、ID (ID) 和收益率 (Return) 的 DataFrame。 CALDT 列转换为 datetime 类型，并添加了年份 (Year) 列。
分组并计算统计量
接下来，我们使用 groupby() 函数按照 "ID" 和 "Year" 进行分组。然后，我们使用 transform() 函数计算每个组的平均收益率和中位数收益率，并将结果乘以 12。transform() 函数会将计算结果广播回原始 DataFrame，保持索引不变。
g = df.groupby(["ID", df.CALDT.dt.year])
return_stats = pd.DataFrame({
"Mean_Return": g["Return"].transform("mean").mul(12),
"Median_Return": g["Return"].transform("median").mul(12)
}).where(g["CALDT"].transform("nunique").ge(2))
df = df.join(return_stats)
print(df)这里，g["Return"].transform("mean").mul(12) 计算了每个组的平均收益率，并乘以 12。 g["Return"].transform("median").mul(12) 计算了每个组的中位数收益率，并乘以 12。.where(g["CALDT"].transform("nunique").ge(2)) 用于筛选组内 "CALDT" 的唯一值数量大于等于 2 的组，不满足条件的组的 "Mean_Return" 和 "Median_Return" 列将被赋值为 NaN。
结果
最终，我们将计算得到的 "Mean_Return" 和 "Median_Return" 列添加到原始 DataFrame 中。
CALDT ID Return Year Mean_Return Median_Return 0 1980-01-31 1 0.02 1980 0.68 0.60 1 1980-02-28 1 0.05 1980 0.68 0.60 2 1980-03-31 1 0.10 1980 0.68 0.60 3 1980-01-31 2 0.05 1980 0.24 0.36 4 1980-02-28 2 -0.02 1980 0.24 0.36 5 1980-03-31 2 0.03 1980 0.24 0.36 6 1980-01-31 3 -0.03 1980 NaN NaN
总结
本文介绍了如何使用 Pandas 的 groupby()、transform() 和 where() 函数，结合条件判断，高效地对 DataFrame 进行分组计算，并将结果应用回原始 DataFrame。这种方法可以灵活地处理各种分组计算需求，并生成包含计算结果的新列。
注意事项：
- transform() 函数的返回值必须与原始 DataFrame 的索引保持一致。
- where() 函数可以根据条件筛选组，并将不满足条件的组的计算结果设置为 NaN。
- 在实际应用中，可以根据具体需求调整分组的依据和计算的统计量。
通过掌握这些技巧，可以更加高效地进行数据分析，并从数据中提取有价值的信息。
