My initial DataFrame looks as follows:
A B quantity
0 1 foo 1
1 1 baz 2
2 1 bar 2
3 1 faz 1
4 2 foo 2
5 2 bar 1
6 3 foo 3
I need to group it by 'A' and make a list of 'B' multiplied by 'quantity':
A B
0 1 [foo, baz, baz, bar, bar, faz]
1 2 [foo, foo, bar]
2 3 [foo, foo, foo]
Currently I'm using groupby() and then apply():
def itemsToList(tdf, column):
collist = []
for row in tdf[column].iteritems():
collist = collist + tdf['quantity'][row[0]]*[row[1]]
return pd.Series({column: collist})
gb = df.groupby('A').apply(itemsToList, 'B')
I doubt it is an efficient way, so I'm looking for a good, "pandaic" method to achieve this.
This could be done in 2 steps, generate a new column that creates the expanded str values, then groupby
on 'A' and apply
list
to this new column:
In [62]:
df['expand'] = df.apply(lambda x: ','.join([x['B']] * x['quantity']), axis=1)
df.groupby('A')['expand'].apply(list)
Out[62]:
A
1 [foo, baz,baz, bar,bar, faz]
2 [foo,foo, bar]
3 [foo,foo,foo]
Name: expand, dtype: object
EDIT
OK after taking inspiration from @Jianxun Li's answer:
In [130]:
df.groupby('A').apply(lambda x: np.repeat(x['B'].values, x['quantity']).tolist())
Out[130]:
A
1 [foo, baz, baz, bar, bar, faz]
2 [foo, foo, bar]
3 [foo, foo, foo]
dtype: object
Also this works:
In [131]:
df.groupby('A').apply(lambda x: list(np.repeat(x['B'].values, x['quantity'])))
Out[131]:
A
1 [foo, baz, baz, bar, bar, faz]
2 [foo, foo, bar]
3 [foo, foo, foo]
dtype: object
Another way to do it. First reshape the df
using pivot_table
and then apply
np.repeat().tolist()
.
import pandas as pd
import numpy as np
df
Out[52]:
A B quantity
0 1 foo 1
1 1 baz 2
2 1 bar 2
3 1 faz 1
4 2 foo 2
5 2 bar 1
6 3 foo 3
df.pivot('A','B','quantity').fillna(0).apply(lambda row: np.repeat(row.index.values, row.values.astype(int)).tolist(), axis=1)
Out[53]:
A
1 [bar, bar, baz, baz, faz, foo]
2 [bar, foo, foo]
3 [foo, foo, foo]
dtype: object
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With