如何在Pandas中实现支持任意分组变量数量的分组求和自定义函数？

你完全可以通过Python的可变位置参数来实现这个泛化版本，不用再为不同数量的分组变量写重复的函数啦。

最终泛化函数 #

def grp_and_sum(data, sum_var, *grp_vars):
    # 将可变参数转为列表传给groupby
    grouped = data.groupby(list(grp_vars))[sum_var].sum()
    return grouped

核心原理 #

这里的*grp_vars是Python的可变位置参数语法——它会把你传入的所有分组变量参数自动打包成一个元组，我们只需要把这个元组转换成列表（pandas的groupby既接受列表也接受元组），就能直接用于分组操作了。

用法示例 #

先用上你提供的测试DataFrame：

import pandas as pd

df = pd.DataFrame({ 
    'cat_1': ['A'] * 3 + ['B'] * 3, 
    'cat_2': ['x', 'y', 'z'] * 2, 
    'value': [1, 2, 3, 4, 5, 6] 
})

1. 单个分组变量（对应你原来的grp_and_sum）

grp_and_sum(df, 'value', 'cat_1')

输出结果：

cat_1
A     6
B    15
Name: value, dtype: int64

2. 两个分组变量（对应你原来的grp_and_sum_2）

grp_and_sum(df, 'value', 'cat_1', 'cat_2')

输出结果：

cat_1  cat_2
A      x        1
       y        2
       z        3
B      x        4
       y        5
       z        6
Name: value, dtype: int64

3. 更多分组变量（如果有需要）

假设你的DataFrame新增了第三列分类列cat_3：

df['cat_3'] = ['foo', 'bar', 'foo'] * 2
grp_and_sum(df, 'value', 'cat_1', 'cat_2', 'cat_3')

同样能正常返回按三个维度分组求和的结果。

参数顺序小说明 #

我把sum_var调整到了*grp_vars前面，这是因为可变参数必须放在位置参数的后面。如果你更习惯原来的参数顺序，也可以改用关键字参数强制指定sum_var，写法如下：

def grp_and_sum(data, *grp_vars, sum_var):
    grouped = data.groupby(list(grp_vars))[sum_var].sum()
    return grouped

# 调用时必须显式指定sum_var
grp_and_sum(df, 'cat_1', sum_var='value')

不过第一种写法调用起来更符合直觉，和你原来的函数使用习惯更接近~

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