How to group by two columns, calculate weighted mean, return DataFrame, in Python

user_7391466 0
Viewed 378

Goal: Take a DataFrame, group by two columns of that DataFrame, calculate the weighted mean of a third column, and return a DataFrame. What adjustments are needed in Reprex below?

##  REPREX
import pandas as pd
import numpy as np

# Turn into DataFrame
md = {"group1": ["A", "A", "B", "B", "A"],
      "group2": [1, 0, 1, 1, 1],
      "var1":[343, 345, 567, 345, 212],
      "wt":[1.2, 1.3, 1.8, .9, 1.9]}

# Turn into DataFrame
my_df = pd.DataFrame(md)

# This works to get unweighted means:
grouped__means = my_df.groupby(["group1", "group2"], as_index=False)[["var1"]].mean()

# This does NOT work for weighted means:
grouped__wtd_means = my_df.groupby(["group1", "group2"], as_index=False)[["var1"]].agg(np.average(weights=my_df["wt"]))

# What needs to change?
pandas dataframe numpy python
2 Answers

You can use GroupBy.apply here.

df.groupby(['group1', 'group2']).apply(lambda x: np.average(x['var1'], weights=x['wt']))

group1  group2
A       0         345.000000
        1         262.709677
B       1         493.000000
dtype: float64
user_124164530

With many groups calculate the weighted average yourself to avoid the apply. Multiply the values by the weights and sum, then divide by the sum of the weights. pandas will align on the Index for the division.

(df['var1'].mul(df['wt']).groupby([df['group1'], df['group2']]).sum()
     .divide(df.groupby(['group1', 'group2'])['wt'].sum()))

group1  group2
A       0         345.000000
        1         262.709677
B       1         493.000000
dtype: float64
user_43333590
Related