Say I have a pandas DataFrame such as:
A B id
0 1 1 0
1 2 1 0
2 3 2 1
3 0 2 1
Say I want to combine rows with the same id
You could use groupby for that with groupby agg
method and tolist method of Pandas Series:
In [762]: df.groupby('id').agg(lambda x: x.tolist())
Out[762]:
A B
id
0 [1, 2] [1, 1]
1 [3, 0] [2, 2]
groupby
return an Dataframe as you want:
In [763]: df1 = df.groupby('id').agg(lambda x: x.tolist())
In [764]: type(df1)
Out[764]: pandas.core.frame.DataFrame
To exactly match your expected result you could additionally do reset_index
or use as_index=False
in groupby
:
In [768]: df.groupby('id', as_index=False).agg(lambda x: x.tolist())
Out[768]:
id A B
0 0 [1, 2] [1, 1]
1 1 [3, 0] [2, 2]
In [771]: df1.reset_index()
Out[771]:
id A B
0 0 [1, 2] [1, 1]
1 1 [3, 0] [2, 2]