CODEWITHSUNDEEP
pythonpandaspandas-groupby
idt_tt
idt_tt

Reputation: 601

Calculating Weighted Average groupby in pandas

I have a dataframe which looks like this:

pd.DataFrame({'A': ['C1', 'C1', 'C1', 'C1', 'C2', 'C2', 'C2'],
                   'W': [0.2, 0.1, 0.5, 0.3, 0.4, 0.3, 0.7],
                   'Y': [2, 0, 4, 3, 2, 2, 3]})

    A    W  Y
0  C1  0.2  2
1  C1  0.1  0
2  C1  0.5  4
3  C1  0.3  3
4  C2  0.4  2
5  C2  0.3  2
6  C2  0.7  3

I would like to calculate the weighted average of each group and add a column with the result.

W = weight

Y = value

resulting dataframe should look like this

    A    W  Y   result
0  C1  0.2  2   3.3
1  C1  0.1  0   3.3
2  C1  0.5  4   3.3
3  C1  0.3  3   3.3
4  C2  0.4  2   3.5
5  C2  0.3  2   3.5
6  C2  0.7  3   3.5

I would like to achieve this with 1 line of code (no function).

Upvotes: 2

Views: 152

Answers (1)

BENY
BENY

Reputation: 323226

Try with transform before get the product of Weight and value

df['result'] = df.Y.mul(df.W).groupby(df.A).transform('sum')
Out[65]: 
0    3.3
1    3.3
2    3.3
3    3.3
4    3.5
5    3.5
6    3.5
dtype: float64

Upvotes: 1

Related Questions