Pythonic way to get a weighted sum using a DataFrame and dict

user_2136 0
Viewed 236

Given a DataFrame and a Dict, how might I get a new column with the weighted sum? The keys in the Dict match the column names in the DataFrame

import pandas as pd
df = pd.DataFrame({'Index': ['aaa', 'bbb', 'ccc'],
                   'a': [1, 2, 3],
                   'b': [4, 5, 6],
                   'c': [7, 8, 9]})
df.set_index('Index', inplace=True)

weights = {'a': 0.5,
        'b': 0.2
       }

Here I want a pythonic way (I'm currently doing a big loop!) to get the following result:

| Index   | a   | b   | c   | weighted_sum            |
| aaa     | 1   | 4   | 7   | (1 * 0.5) + (4 * 0.2)   |
| bbb     | 2   | 5   | 8   | (2 * 0.5) + (5 * 0.2)   | 
| ccc     | 3   | 6   | 9   | (3 * 0.5) + (6 * 0.2)   |

Note, I won't always know what keys are in the weights dict

pandas python
4 Answers

You can do:

df['weighted_sum'] = df[['a','b']].mul(weights).sum(1)

Or equivalently:

df['weighted_sum'] = (df[['a','b']] * weights).sum(1)

Output:

       a  b  c  weighted_sum
Index                       
aaa    1  4  7           1.3
bbb    2  5  8           2.0
ccc    3  6  9           2.7

Now if you don't always know the keys, you can get the keys with weights.keys():

(df[weights.keys()] * weights).sum(1)

That might throw a KeyError if your keys are not in the dataframe's columns. In which case, you can just turn weights into a series and Pandas will do the heavy lifting for you.

df.mul(pd.Series(weights)).sum(1)
user_42384080

You can use .assign then .mul and .sum (dont use dict as variable name)

import pandas as pd
df = pd.DataFrame({'Index': ['aaa', 'bbb', 'ccc'],
                   'a': [1, 2, 3],
                   'b': [4, 5, 6],
                   'c': [7, 8, 9]})
df.set_index('Index', inplace=True)

pdict = {'a': 0.5,
        'b': 0.2
       }
       
df = df.assign(**pdict).mul(df)[['a','b']].sum(1)

print(df)
user_122698570

a simple way to do it would be like this

def funk(df,d):
    return df[list(d)[0]]*d[list(d)[0]] + df[list(d)[1]]*d[list(d)[1]]

funk(df,d)

results

aaa    1.3
bbb    2.0
ccc    2.7
dtype: float64
user_143434650

My one-line version looks like:

df['weighted_sum'] = sum([df[key] * weights[key] for key in weights.keys()])

Printing results in

       a  b  c  weighted_sum
Index                       
aaa    1  4  7           1.3
bbb    2  5  8           2.0
ccc    3  6  9           2.7
user_44682460
Related