Groupby с пользовательскими функциями Pandas

Я понимаю, что передача функции в качестве группового ключа вызывает функцию один раз для каждого значения индекса, а возвращаемые значения используются в качестве имен групп. Что я не могу понять, так это как вызывать функцию для значений столбцов.

Так что я могу сделать это:

people = pd.DataFrame(np.random.randn(5, 5), 
                      columns=['a', 'b', 'c', 'd', 'e'],
                      index=['Joe', 'Steve', 'Wes', 'Jim', 'Travis'])
def GroupFunc(x):
    if len(x) > 3:
        return 'Group1'
    else:
        return 'Group2'

people.groupby(GroupFunc).sum()

Это разбивает данные на две группы, одна из которых имеет значения индекса длиной 3 или менее, а другая - длиной три или более. Но как я могу передать одно из значений столбца? Так, например, если значение столбца d для каждой точки индекса больше 1. Я понимаю, что мог бы просто сделать следующее:

people.groupby(people.a > 1).sum()

Но я хочу знать, как сделать это в пользовательской функции для дальнейшего использования.

Что-то вроде:

def GroupColFunc(x):
if x > 1:
    return 'Group1'
else:
    return 'Group2'

Но как я это называю? Я старался

people.groupby(GroupColFunc(people.a))

и подобные варианты, но это не работает.

Как передать значения столбца в функцию? Как бы я передал несколько значений столбцов, например, для группы, например, people.a> people.b?

Ответ 1

Чтобы сгруппировать по > 1, вы можете определить свою функцию следующим образом:

>>> def GroupColFunc(df, ind, col):
...     if df[col].loc[ind] > 1:
...         return 'Group1'
...     else:
...         return 'Group2'
...

Затем вызывается как

>>> people.groupby(lambda x: GroupColFunc(people, x, 'a')).sum()
               a         b         c         d        e
Group2 -2.384614 -0.762208  3.359299 -1.574938 -2.65963

Или вы можете сделать это только с анонимной функцией:

>>> people.groupby(lambda x: 'Group1' if people['b'].loc[x] > people['a'].loc[x] else 'Group2').sum()
               a         b         c         d         e
Group1 -3.280319 -0.007196  1.525356  0.324154 -1.002439
Group2  0.895705 -0.755012  1.833943 -1.899092 -1.657191

Как сказано в документации, вы также можете группировать, передавая Series, предоставляя сопоставление имен меток → :

>>> mapping = np.where(people['b'] > people['a'], 'Group1', 'Group2')
>>> mapping
Joe       Group2
Steve     Group1
Wes       Group2
Jim       Group1
Travis    Group1
dtype: string48
>>> people.groupby(mapping).sum()
               a         b         c         d         e
Group1 -3.280319 -0.007196  1.525356  0.324154 -1.002439
Group2  0.895705 -0.755012  1.833943 -1.899092 -1.657191