Pandas数据分组条件计算与结果回填教程

在数据分析工作中，我们经常需要对数据集进行分组操作，并对每个分组内的特定列执行聚合计算。然而，更复杂的场景可能要求我们根据分组的某些属性（例如，分组内记录的数量是否满足某个阈值）来决定是否应用这些计算，并将结果“广播”回原始dataframe的每一行，而非仅仅得到一个聚合后的摘要。本教程将详细阐述如何在pandas中高效地实现这一需求。

数据准备

首先，我们创建一个示例DataFrame，其中包含交易日期（CALDT）、唯一标识符（ID）和回报率（Return）。为了后续按年份分组，我们还会从CALDT列中提取年份信息。

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]
     })

# 将CALDT转换为日期时间类型并提取年份
df['CALDT'] = pd.to_datetime(df['CALDT'])
df['Year'] = df['CALDT'].dt.year

print("原始DataFrame:")
print(df)

核心解决方案：条件分组计算与结果回填

我们的目标是：对于每个ID和年份的组合，如果该ID在该年份内“存活”至少2个月（即有至少2条记录），则计算其Return的年化均值和中位数（乘以12），并将这些值回填到该分组的每一行；否则，对应的结果应为NaN。

Pandas提供了groupby()结合transform()和where()的强大组合来解决这类问题。

1. 分组操作

首先，我们需要根据ID和Year对DataFrame进行分组。

g = df.groupby(["ID", "Year"])

这里，g是一个DataFrameGroupBy对象，它包含了分组的信息，但尚未执行任何计算。

2. 使用 transform() 进行组内计算并广播结果

transform()方法是实现将聚合结果“广播”回原始DataFrame的关键。与agg()方法不同，transform()返回的结果具有与原始DataFrame相同（或可对齐）的索引和形状，这使得我们可以直接将其与原始DataFrame合并。

我们计算Return的均值和中位数，并将其乘以12进行年化。

return_stats = pd.DataFrame({
    "Mean_Return": g["Return"].transform("mean").mul(12),
    "Median_Return": g["Return"].transform("median").mul(12)
})

此时，return_stats DataFrame已经包含了每个分组的年化均值和中位数，并且这些值被重复填充到该分组的所有行中。

算家云

高效、便捷的人工智能算力服务平台

37

查看详情

3. 应用条件筛选 where()

接下来，我们需要应用“如果ID在该年份内至少存活2个月”的条件。我们可以通过计算每个分组中CALDT列的唯一值数量来实现这一点。transform("nunique")可以计算每个分组中CALDT的唯一值数量，并将其广播回原始DataFrame的形状。然后，我们使用ge(2)（greater than or equal to 2）来生成一个布尔掩码。

where()方法允许我们根据条件选择性地保留或替换DataFrame中的值。如果条件为False，则where()会将对应位置的值替换为NaN（默认行为）。

# 计算每个分组的唯一月份数，并检查是否大于等于2
condition = g["CALDT"].transform("nunique").ge(2)

# 根据条件应用筛选
return_stats = return_stats.where(condition)

现在，return_stats中不满足条件（即唯一月份数小于2）的行已经被替换为NaN。

4. 合并结果

最后一步是将计算出的统计量合并回原始的DataFrame df。由于return_stats的索引与df的索引是匹配的，我们可以使用join()方法。

df_final = df.join(return_stats)

print("\n最终结果DataFrame:")
print(df_final)

完整代码示例

将上述步骤整合到一起，得到完整的解决方案代码：

import pandas as pd
import numpy as np

# 1. 数据准备
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

# 2. 核心解决方案
# 2.1. 分组
g = df.groupby(["ID", "Year"])

# 2.2. 计算统计量并广播
return_stats = pd.DataFrame({
    "Mean_Return": g["Return"].transform("mean").mul(12),
    "Median_Return": g["Return"].transform("median").mul(12)
})

# 2.3. 应用条件筛选
# 计算每个分组的唯一月份数，并检查是否大于等于2
condition = g["CALDT"].transform("nunique").ge(2)
return_stats = return_stats.where(condition)

# 2.4. 合并结果
df_final = df.join(return_stats)

print(df_final)

关键方法解析

• df.groupby(["ID", "Year"]): 这是Pandas中进行分组操作的基础。它根据指定的列（ID和Year）将DataFrame划分为多个子组。
• .transform("function_name"): transform方法在每个分组上应用一个函数，并将结果返回一个与原始DataFrame具有相同索引和形状的Series或DataFrame。这意味着每个分组的聚合结果会被“广播”回该分组的所有原始行。这与agg或apply在返回聚合结果时的行为不同。
  • 例如，g["Return"].transform("mean")会计算每个ID和Year分组的Return均值，然后将这个均值值填充到该分组的每一行。
• .dt.year: 这是一个日期时间访问器，用于从日期时间序列中提取年份。
• .nunique(): 这个函数用于计算Series中唯一值的数量。在transform("nunique")的上下文中，它会计算每个分组中指定列的唯一值数量。
• .ge(value): 这是一个比较运算符，表示“大于等于”（Greater than or Equal to）。例如，series.ge(2)会返回一个布尔Series，其中每个元素指示对应位置的值是否大于或等于2。
• .where(condition): where方法根据布尔条件选择性地替换DataFrame中的值。如果condition中的对应位置为False，则DataFrame中的该值将被替换为NaN（默认行为）；如果为True，则值保持不变。

注意事项与总结

• 性能优化: 对于大型数据集，transform()通常比结合apply()和自定义函数更高效，因为它在C语言级别实现了许多常见的聚合操作。
• 灵活性: 如果你需要执行更复杂的、transform()不支持的自定义操作，可以考虑使用groupby().apply()。但请注意，apply()的性能开销通常会更高。
• 索引对齐: join()方法在这里能够顺利工作，是因为return_stats的索引与df的索引是完全对齐的。在进行此类操作时，确保索引的一致性是至关重要的。

通过掌握groupby()、transform()和where()的组合使用，你可以在Pandas中高效地执行复杂的条件分组计算，并将结果无缝地集成回原始数据集，从而极大地提升数据处理的灵活性和效率。

以上就是Pandas数据分组条件计算与结果回填教程的详细内容，更多请关注php中文网其它相关文章！