Weighted Mean row wise Pandas

Viewed 71

I would like to compute the weighted mean for each row, for a given weight 'wt':

df = pd.DataFrame(  {
'datetime': ['2015-01-02', '2015-01-03', '2015-01-04', '2015-01-05', '2015-01-06'],
'var1': [0, 0, 1, 1, 1],
'var2': [-1, -2, 1, 2, 1]})
df['datetime'] = pd.to_datetime(df['datetime'])
df.set_index([ 'datetime'], inplace =True)
wt = [0,1]

which looks like:

     datetime  var1  var2
0  2015-01-02     0    -1
1  2015-01-03     0    -2
2  2015-01-04     1     1
3  2015-01-05     1     2
4  2015-01-06     1     1

I want to add "weighted_mean" column:

df['weighted_mean'] = [-1, -2, 1, 2, 1]

Expected output:

            var1  var2  weighted_mean
datetime                             
2015-01-02     0    -1             -1
2015-01-03     0    -2             -2
2015-01-04     1     1              1
2015-01-05     1     2              2
2015-01-06     1     1              1
1 Answers

IIUC, you can multiply the relevant columns with wt and sum row-wise:

df['weighted_mean'] = df[['var1', 'var2']].mul(wt).sum(axis=1)

Output:

            var1  var2  weighted_mean
datetime                             
2015-01-02     0    -1             -1
2015-01-03     0    -2             -2
2015-01-04     1     1              1
2015-01-05     1     2              2
2015-01-06     1     1              1
Related