Groupby transform to list in pandas does not work

Viewed 1079

Best described with an example

import pandas as pd

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

And i want to create a column that contains in a list the elements of column b by group of column a

resulting in the following

   a    b      c
0  A    1  [1, 4, 7]
1  A    4  [1, 4, 7]
2  A    7  [1, 4, 7]
3  B    2  [2, 5, 8]
4  B    5  [2, 5, 8]
5  B    8  [2, 5, 8]
6  C    3  [3, 6, 9]
7  C    6  [3, 6, 9]
8  C    9  [3, 6, 9]

I can do this with groupby and apply or agg and then joining the dataframes like so

df_tmp = df.groupby('a')['b'].agg(list).reset_index()
df.merge(df_tmp, on='a')

But i would also be expecting to do the same with transform

df['c'] = df.groupby('a')['b'].transform(list)

but the column c is the same as column b

Also the following

df.groupby('a')['b'].transform(lambda x: len(x))

return a series with the values 3 i.e. the length of the grouped elements is 3 (to be expected)

Also this

df.groupby('a')['b'].transform(lambda x: list(x))

does not provide the expected result.

So to my question, how can i obtain the desired result with groupby and tranform

pandas version is 1.0.5

2 Answers

I come up one fix with below. PS : it should something wrong with transform , when the object type is list tuple or set..

df.groupby('a')['b'].transform(lambda x : [x.tolist()]*len(x))
Out[226]: 
0    [1, 4, 7]
1    [1, 4, 7]
2    [1, 4, 7]
3    [2, 5, 8]
4    [2, 5, 8]
5    [2, 5, 8]
6    [3, 6, 9]
7    [3, 6, 9]
8    [3, 6, 9]
Name: b, dtype: object

Interesting problem, not sure what happens with transform in the background. One go-around is to map with groupby().agg():

df['c'] = df['a'].map(df.groupby('a')['b'].agg(list))

Output:

   a  b          c
0  A  1  [1, 4, 7]
1  B  2  [2, 5, 8]
2  C  3  [3, 6, 9]
3  A  4  [1, 4, 7]
4  B  5  [2, 5, 8]
5  C  6  [3, 6, 9]
6  A  7  [1, 4, 7]
7  B  8  [2, 5, 8]
8  C  9  [3, 6, 9]
Related