Abhishek Thakur Abhishek Thakur - 2 months ago 10
Python Question

grouping rows in list in pandas groupby

I have a pandas data frame like:

A 1
A 2
B 5
B 5
B 4
C 6


I want to group by the first column and get 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

You can do this using groupby to group on the column of interest and then apply list to every group:

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

[6 rows x 2 columns]

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

Out[76]:
a
A       [1, 2]
B    [5, 5, 4]
C          [6]
Name: b, dtype: object