如何在Pandas中计算条件响应概率？

基础数据与初始聚合 #

先看示例数据构造代码：

import pandas as pd

gender = [0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1]
is_family = [0,0,0,0,1,1,1,1,0,0,0,0,1,1,1,1]
treatment = [0,1,0,1,0,1,0,1,0,1,0,1,0,1,0,1]
response = [1,0,0,1,1,0,0,1,1,0,0,1,1,0,0,1]
num_rows = [10,10,5,20,0,5,10,30,20,30,10,5,60,10,10,20]

df = pd.DataFrame(data={'gender': gender, 'is_family': is_family, 'treatment': treatment, 'response': response, 'num_rows': num_rows})

你已经实现了按gender、treatment、response分组求和num_rows的代码：

agg_df = df.groupby(by=['gender', 'treatment', 'response'])['num_rows'].sum().reset_index()

计算条件响应概率 #

要得到给定gender和treatment组合下的response概率，核心逻辑是用每组的num_rows总和，除以对应gender+treatment分组的总num_rows。以下是两种简洁的实现方式：

方法1：使用transform直接计算

这种方法可以在聚合后的DataFrame上直接生成概率列，步骤更紧凑：

# 先执行基础聚合
agg_df = df.groupby(by=['gender', 'treatment', 'response'])['num_rows'].sum().reset_index()
# 按gender+treatment分组计算总num_rows，再用当前行的num_rows除以该总和得到概率
agg_df['resp_prob'] = agg_df['num_rows'] / agg_df.groupby(['gender', 'treatment'])['num_rows'].transform('sum')

方法2：分步计算再合并

如果需要更清晰的流程拆分，可以先单独计算分组总规模，再通过合并关联计算概率：

# 步骤1：完成基础聚合
agg_df = df.groupby(by=['gender', 'treatment', 'response'])['num_rows'].sum().reset_index()
# 步骤2：计算每个gender+treatment组合的总num_rows
total_df = agg_df.groupby(['gender', 'treatment'])['num_rows'].sum().reset_index(name='total_rows')
# 步骤3：合并数据并计算概率
result_df = pd.merge(agg_df, total_df, on=['gender', 'treatment'])
result_df['resp_prob'] = result_df['num_rows'] / result_df['total_rows']
# 可选：删除中间辅助列
result_df = result_df.drop('total_rows', axis=1)

最终结果 #

运行上述任意一种方法后，都会得到符合预期的结果：

gender  treatment  response  num_rows  resp_prob
0       0          0         0        15   0.600000
1       0          0         1        10   0.400000
2       0          1         0        15   0.230769
3       0          1         1        50   0.769231
4       1          0         0        20   0.200000
5       1          0         1        80   0.800000
6       1          1         0        40   0.615385
7       1          1         1        25   0.384615

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