Odtnhj

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?

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

If performance is important go down to numpy level:

import numpy as np

df = pd.DataFrame('a': np.random.randint(0,60,600), 'b': [1,2,5,5,4,6]*100)

def f(df):
 keys, values = df.sort_values('a').values.T
 ukeys, index = np.unique(keys,True)
 arrays = np.split(values,index[1:])
 df2 = pd.DataFrame('a':ukeys,'b':[list(a) for a in arrays])
 return df2

Tests:

In [301]: %timeit f(df)
1000 loops, best of 3: 1.64 ms per loop

In [302]: %timeit df.groupby('a')['b'].apply(list)
100 loops, best of 3: 5.26 ms per loop

As you were saying the groupby method of a pd.DataFrame object can do the job.

Example

 L = ['A','A','B','B','B','C']
 N = [1,2,5,5,4,6]

 import pandas as pd
 df = pd.DataFrame(zip(L,N),columns = list('LN'))


 groups = df.groupby(df.L)

 groups.groups
 'A': [0, 1], 'B': [2, 3, 4], 'C': [5]

which gives and index-wise description of the groups.

To get elements of single groups, you can do, for instance

 groups.get_group('A')

 L N
 0 A 1
 1 A 2

 groups.get_group('B')

 L N
 2 B 5
 3 B 5
 4 B 4

A handy way to achieve this would be:

df.groupby('a').agg('b':lambda x: list(x))

Look into writing Custom Aggregations: https://www.kaggle.com/akshaysehgal/how-to-group-by-aggregate-using-py

To solve this for several columns of a dataframe:

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

In [6]: df
Out[6]: 
 a b c
0 A 1 3
1 A 2 3
2 B 5 3
3 B 5 4
4 B 4 4
5 C 6 4

In [7]: df.groupby('a').agg(lambda x: list(x))
Out[7]: 
 b c
a 
A [1, 2] [3, 3]
B [5, 5, 4] [3, 4, 4]
C [6] [4]

This answer was inspired from Anamika Modi's answer. Thank you!