Python pandas groupby分组计算优化及多变量批量执行方案咨询

你可以将统计逻辑封装为接收分组列名作为参数的通用函数，遍历变量列表时逐个传入列名调用即可，具体实现如下：

import pandas as pd

def groupby_stat(df, group_col):
    # 入参group_col为需要和Y共同分组的维度列，如X/U/V/W
    # 第一步：执行分组统计
    group_res = df.groupby([group_col, 'Y']).agg(
        total_stocks=('Stock', 'count'),
        mean_number=('Number', 'mean')
    ).reset_index()
    # pyspark环境请保留下方persist语句，pandas环境直接注释即可
    # group_res = group_res.persist()

    # 第二步：计算占比
    # 注意原代码此处逻辑有误：不能从原始df取total_stocks字段，需从统计结果中分组求和
    group_res['stock_sum'] = group_res.groupby('Y')['total_stocks'].transform('sum')
    group_res['proportion'] = group_res['total_stocks'] / group_res['stock_sum']

    # 可选操作：新增维度标识列，方便后续合并多个维度的结果
    group_res['dim_type'] = group_col
    return group_res

调用方式 #

# 构造需要遍历的维度列表
dim_list = ['X', 'U', 'V', 'W']

# 方案1：所有维度的结果合并为一个总表
all_stat_result = pd.concat(
    [groupby_stat(df, dim) for dim in dim_list],
    ignore_index=True
)

# 方案2：各维度结果单独存储，按需取用
result_dict = {}
for dim in dim_list:
    result_dict[dim] = groupby_stat(df, dim)
# 取用X维度结果直接调用 result_dict['X'] 即可

注意事项 #

• 如果你的数据集列名是全小写格式，请将代码中的列名对应修改为x/y/stock/number即可
• 若使用PySpark框架，仅需调整聚合、transform的对应语法，函数封装的整体逻辑通用

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