本文介绍了如何使用 Pandas 在 DataFrame 分组后，针对每个分组的行应用自定义函数计算特定值。重点在于利用 groupby() 和 transform() 方法，结合条件判断，实现对满足特定条件的分组进行计算，并将结果广播回原始 DataFrame。通过本文，你将掌握一种高效处理分组数据的技巧，并能灵活应用于各种数据分析场景。
Pandas 提供了强大的分组（Grouping）功能，允许你根据一列或多列的值将 DataFrame 拆分成多个组。在许多数据分析场景中，我们需要在每个组内进行计算，并将结果应用回原始 DataFrame 的每一行。本文将详细介绍如何使用 Pandas 的 groupby() 和 transform() 方法，结合条件判断，实现按行应用自定义函数计算特定值。
使用 groupby() 和 transform() 进行分组计算
groupby() 方法用于将 DataFrame 按照指定的列进行分组。transform() 方法则允许你对每个分组应用一个函数，并将结果广播回原始 DataFrame。这与 agg() 方法不同，agg() 方法会返回聚合后的结果，而 transform() 方法会返回与原始 DataFrame 相同大小的结果。
以下是一个示例，演示如何使用 groupby() 和 transform() 计算每个 ID 和年份组合的平均回报率和中位数回报率，并将结果乘以 12。
import pandas as pd
import numpy as np
# 创建示例 DataFrame
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'])
# 按照 ID 和年份进行分组
g = df.groupby(["ID", df.CALDT.dt.year])
# 计算平均回报率和中位数回报率，并乘以 12
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))
# 将计算结果与原始 DataFrame 合并
df = df.join(return_stats)
print(df)代码解释：
- 创建 DataFrame: 首先，我们创建一个包含日期 (CALDT)、ID (ID) 和回报率 (Return) 的 DataFrame。
- 转换为 datetime: 将 CALDT 列转换为 datetime 类型。
- 分组: 使用 groupby(["ID", df.CALDT.dt.year]) 按照 ID 和 CALDT 的年份进行分组。
- 计算统计量: 使用 transform() 方法计算每个分组的平均回报率 (Mean_Return) 和中位数回报率 (Median_Return)，并将结果乘以 12。transform() 方法会将计算结果广播回原始 DataFrame 的每一行。
- 条件判断: 使用 where() 方法，结合 g["CALDT"].transform("nunique").ge(2)，判断每个分组中唯一日期数量是否大于等于 2。如果不满足条件，则将对应的 Mean_Return 和 Median_Return 设置为 NaN。
- 合并结果: 使用 join() 方法将计算得到的 return_stats DataFrame 与原始 DataFrame 合并。
注意事项
- 确保分组列的数据类型正确。在上面的例子中，我们需要将 CALDT 列转换为 datetime 类型，才能正确提取年份。
- transform() 方法返回的结果必须与原始 DataFrame 的大小相同。
- where() 方法用于根据条件判断是否保留计算结果。如果条件不满足，则将结果设置为 NaN。
总结
本文介绍了如何使用 Pandas 的 groupby() 和 transform() 方法，结合条件判断，实现按行应用自定义函数计算特定值。这种方法可以高效地处理分组数据，并能灵活应用于各种数据分析场景。通过掌握这种技巧，你可以更轻松地进行复杂的数据处理和分析。
