Pandas dataframe how to make element-wise mean of list columns

user_6057371 0
Viewed 151

I have a dataframe with categorical features and values features. Value features are always lists with the exact same size (let's say 3 for the example). So the dataframe is:

Sub Model iter key     l1       l2     l3
 01  b     0    0    [8,8,5] [3,8,1] [6,7,8]
 01  b     1    1    [4,3,6] [3,4,0] [4,0,4]
 01  b     2    2    [0,1,4] [0,0,5] [8,2,3]
 03  b     3    0    [1,8,2] [4,6,0] [1,3,9]
 03  b     4    1    [7,3,1] [6,8,1] [6,7,9]
 03  b     5    2    [1,1,0] [11,4,8] [8,5,9]

I want to group the dataframe by [sub,model], such that in each row I will take the mean over the values of the columns key. So I will get:

Sub  Model     l1         l2       l3 
 01   b      [4, 4, 5]  [2,2,2]   [6,3,5]
 03   b      [3, 4, 1]  [7,6,3]   [3,5,9]

What will be the best way to do so?

data-munging pandas-groupby pandas dataframe python
3 Answers

You can use agg and a custom function:

def mean_list(sr):
    return sr.apply(pd.Series).mean()

out = df.groupby(['Sub', 'Model'])[['l1', 'l2', 'l3']].agg(mean_list)

>>> out
                        l1               l2               l3
Sub Model
01  b      [4.0, 4.0, 5.0]  [2.0, 4.0, 2.0]  [6.0, 3.0, 5.0]
03  b      [3.0, 4.0, 1.0]  [7.0, 6.0, 3.0]  [5.0, 5.0, 9.0]
user_152399510

You can use the following snippet:

(df.groupby(['Sub', 'Model'])
   [['l1', 'l2', 'l3']]
   .apply(lambda d: pd.Series({c: np.array([*d[c].values]).mean(0)
                               for c in d.columns}))
)

output:

                        l1               l2               l3
Sub Model                                                   
01  b      [4.0, 4.0, 5.0]  [2.0, 4.0, 2.0]  [6.0, 3.0, 5.0]
03  b      [3.0, 4.0, 1.0]  [7.0, 6.0, 3.0]  [5.0, 5.0, 9.0]
user_163434640

You can use .groupby() + .agg() together with np.mean(), as follows:

import numpy as np

(df.groupby(['Sub', 'Model'])[['l1', 'l2', 'l3']]
   .agg(lambda x: np.mean(x.values.tolist(), axis=0)).reset_index()
)

Result:

  Sub Model               l1               l2               l3
0  01     b  [4.0, 4.0, 5.0]  [2.0, 4.0, 2.0]  [6.0, 3.0, 5.0]
1  03     b  [3.0, 4.0, 1.0]  [7.0, 6.0, 3.0]  [5.0, 5.0, 9.0]

Edit:

If your l1, l2, l3 are actually strings look like lists, you can do the following before the solution codes:

df['l1'] = df['l1'].str.strip('[]').str.split(',', expand=True).astype(int).apply(lambda x: list(x), axis=1)
df['l2'] = df['l2'].str.strip('[]').str.split(',', expand=True).astype(int).apply(lambda x: list(x), axis=1)
df['l3'] = df['l3'].str.strip('[]').str.split(',', expand=True).astype(int).apply(lambda x: list(x), axis=1)
user_150706970
Related