Question

我的数据如下：

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

我想创建一个分组（具有可变数目的分组）并汇总我的数据的函数。例如，以下功能分别针对一组和两组实现此目的。

def grp_and_sum(data, grp_var, sum_var):

    df = data.groupby([grp_var])[sum_var]\
         .sum()

    return(df)

def grp_and_sum_2(data, grp_var1, grp_var2, sum_var):

    df = data.groupby([grp_var1, grp_var2])[sum_var]\
        .sum()

    return(df)

除了可变数量的“分组”变量外，这些函数几乎相同。如何泛化第一个函数以接受任意数量的分组变量？谢谢。

Answer 1

您可以将varargs用于石斑鱼，但是您的sum_var参数必须作为关键字参数传递。

def grp_and_sum_n(data, *args, sum_var):
    return data.groupby([*args])[sum_var].sum()

grp_and_sum_n(df, 'cat_2', sum_var='value')
cat_2
x    5
y    7
z    9
Name: value, dtype: int64

grp_and_sum_n(df, 'cat_1', 'cat_2', sum_var='value')
cat_1  cat_2
A      x        1
       y        2
       z        3
B      x        4
       y        5
       z        6
Name: value, dtype: int64

Answer 2

我选择删除[]附近的grp_var。这为groupby提供了by参数的完全灵活性：

依据：映射，功能，标签或标签列表

def grp_and_sum(data, by, sum_var):
    df = data.groupby(by)[sum_var].sum()
    return(df)

grp_and_sum(df, 'cat_1', sum_var='value')
#cat_1
#A     6
#B    15
#Name: value, dtype: int64

grp_and_sum(df, ['cat_1'], sum_var='value')
#cat_1
#A     6
#B    15
#Name: value, dtype: int64

grp_and_sum(df, ['cat_1', 'cat_2'], sum_var='value')
#cat_1  cat_2
#A      x        1
#       y        2
#       z        3
#B      x        4
#       y        5
#       z        6
#Name: value, dtype: int64

甚至可以为索引传递函数

grp_and_sum(df, lambda x: x//2, sum_var='value')
#0     3
#1     7
#2    11
#Name: value, dtype: int64

用户定义的函数，在pandas中具有任意数量的分组变量

2 个答案: