CODEWITHSUNDEEP
pythonpandasdataframeweighted-average
user2590177
user2590177

Reputation: 177

Weighted average for each row of a pandas dataframe

We have a dataframe df defined as such:

t = pd.DataFrame(
{
    "id": ["id1", "id2", "id3", "id4"],
    "A": [1, 4, 6, 12],
    "B": [5, 8, 3, 6],
    "C": [9, 14, 7, 10],
}
)

Then I have a list:

weight = [2, .5, 1]

I need to create a new column "WMean" giving for each row the weighted average where column A has a weight 2, column B a weight .5, and column C a weight 1. Weight does not have to be a list. It can have another type.

I would need the fastest way to do it.

The result is:

id  |  A | B |  C | WMean
id1 |  1 | 5 |  9 |  3.857142857142857
id2 |  4 | 8 | 14 |  7.428571428571429
id3 |  6 | 3 |  7 |  5.857142857142857
id4 | 12 | 6 | 10 | 10.571428571428571

Thanks.

Upvotes: 1

Views: 3921

Answers (1)

jezrael
jezrael

Reputation: 863751

Use numpy.average with filtered columns by list:

weight = [2, .5, 1]
cols = ['A','B','C']

t['WMean'] = np.average(t[cols], weights=weight, axis=1)
print (t)
    id   A  B   C      WMean
0  id1   1  5   9   3.857143
1  id2   4  8  14   7.428571
2  id3   6  3   7   5.857143
3  id4  12  6  10  10.571429

Upvotes: 8

Related Questions