grouping rows in list in pandas groupby

Question

I have a pandas data frame like:

a b 

A 1 

A 2     

B 5 

B 5 

B 4 

C 6

I want to group by the first column and get the second column as lists in rows:

A [1,2] 

B [5,5,4] 

C [6]

Is it possible to do something like this using pandas groupby?

Answer 1

You can do this using pandas groupby function. Pandas apply function takes another function as an argument to manipulate the selected dataframe  :

For example:

In [1]: # create the dataframe

df = pd.DataFrame( {'a':['A','A','B','B','B','C'], 'b':[1,2,5,5,4,6]})

df 

Out[1]:

   a b 

0 A 1 

1 A 2 

2 B 5

3 B 5 

4 B 4 

5 C 6 

In [76]: df.groupby('a')['b'].apply(list) 

Out[76]:

a 

A [1, 2] 

B [5, 5, 4] 

C [6]