Python - кадр данных Pandas: настраиваемая функция агрегации после группы?

Question

Я использую некоторые функции агрегации после выполнения GroupBy на панд dataframe, как:

my_df.groupby(['id']).agg(['count']) 

Я интересно можно ли иметь настроенную функцию агрегации? Например, в моем кадре данных:

id  color 
--------------------  
001  red 
001  blue 
001  yellow 
002  green 
002  black 
003  yellow 
003  white 
003  blue 

Я хочу создать пользовательскую функцию с именем all_color, так что я мог бы сделать что-то вроде:

my_df.groupby(['id']).agg(['all_color']) 

и получить кадр выходных данных, как:

id  all_color 
-------------------- 
001  [red,blue,yellow] 
002  [green,black] 
003  [yellow,white,blue] 

Answer 1

Использовать функцию apply, а метод tolist() для преобразования в значения в список.

In [12]: df.groupby('id')['color'].apply(lambda x: x.tolist()) 
Out[12]: 
id 
1  [red, blue, yellow] 
2   [green, black] 
3 [yellow, white, blue] 
Name: color, dtype: object 

Использование reset_index для преобразования серии в dataframe

In [21]: df.groupby('id')['color'].apply(lambda x: x.tolist()).reset_index() 
Out[21]: 
    id     color 
0 1 [red, blue, yellow] 
1 2   [green, black] 
2 3 [yellow, white, blue] 

Answer 2

Если вы хотите, чтобы это как dataframe вы можете использовать pivot_table:

In [11]: pd.pivot_table(df, values="id", index=df["id"], columns=df["color"], aggfunc='count', fill_value=0) 
Out[11]: 
color black blue green red white yellow 
id 
1   0  1  0 1  0  1 
2   1  0  1 0  0  0 
3   0  1  0 0  1  1 

Примечание: это очень похоже на выход от get_dummies.