Пользовательская функция с произвольным числом группирующих переменных в пандах

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.Это дает вам полную гибкость аргумента by для groupby:

путем: отображения, функции, метки или списка меток

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