How to calculate weighted mean using group_by function in Pandas

user_12228905 0
Viewed 73

I have a sample data frame that looks like below (4*43). I would like to group_by "id" and then calculate the weighted mean for each column (0,1,2,...40) by "length". How should I go about combining group_by with the weighted mean calculation?

id length 0    1      2    3 ....   40
1   10    0.1  0.02   0.7  0.1 ..  0.4
1   5     0.2  0.01   0.1  0.4 ... 0.1
3   6     0.1  0.05.  0.2  0.5 ... 0.05
3   5     0.5  0.7    0.6  0.5 ... 0.8

I tried to use the group by function as below to get just the grouped average, but it failed to work.

df.groupby("id", as_index=True).mean().iloc[3:43]

Ideal output:

   id  weighted_0    weighted_1      weighted_2  ....   weighted_40
    1    
    3
pandas dataframe python
2 Answers

Another option without using apply (which is generally not recommended for performance reasons, see timing example below):

(df.iloc[:, 2:]
    .multiply(df['length'], axis=0)
    .divide(df.groupby('id')['length'].transform('sum'), axis=0)
    .groupby(df['id'])
    .sum()
    .add_prefix('weighted_'))

Output:

    weighted_0  weighted_1  weighted_2  weighted_3  weighted_40
id                                                             
1     0.133333    0.016667    0.500000         0.2     0.300000
3     0.281818    0.345455    0.381818         0.5     0.390909

For example, timing on a 10K rows dataset with 5K unique ids:

%timeit f_this_answer(df)
2.87 ms ± 18.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit f_with_apply(df)
2.04 s ± 5.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

This answer runs in 3 ms compared with 2 s execution time for Jérôme Richard's answer with apply (approximately 700x improvement)

user_67927430

You can mix apply of Pandas with the weighted average of Numpy like this:

colNames = [f'weighted_{i}' for i in range(len(df.columns)-2)]

def weightedMeans(subDf):
    tmp = np.average(subDf.iloc[:,2:], axis=0, weights=subDf['length'])
    return pd.Series(tmp, index=colNames)

df.groupby('id').apply(weightedMeans).reset_index()

Here is the result:

   id  weighted_0  weighted_1  weighted_2  weighted_3
0   1    0.133333    0.016667    0.500000         0.2
1   3    0.281818    0.345455    0.381818         0.5
user_129395570
Related