How to match column names with dictionary keys and add value to counter

Question

I created a dataframe that has binary values for each cell, where each row is a user and each column is a company the user can select (or not), like this:

company1 company2 company3
1        0        0
0        0        1
0        1        1

And I created a dictionary that categorizes each company into either a high, mid, or low value company:

{'company1': 'high',
'company2': 'low',
'company3': 'low'}

Currently there are companies that are in the dataframe but not in the dictionary, but this should be fixed relatively soon. I would like to create variables for how many times each user selected a high, mid, or low value company. Ultimately should look something like this:

company1 company2 company3 total_low total_mid total_high
1        0        0         0         0         1
0        0        1         1         0         0
0        1        1         2         0         0

I started creating a loop to accomplish this, but I'm not sure how to match the column name with the dictionary key/value, or if this is even the most efficient method (there are ~18,000 rows/users and ~100 columns/companies in total):

total_high = []
total_mid = []
total_low = []
for row in range(df.shape[0]):
    for col in range(df.shape[1]):
        if df.iloc[row,col] == 1:
            # match column name with dict key and add value to 
            # counter

user_4238408 · Answer

One possible approach:

d = {'company1': 'high',
     'company2': 'low',
     'company3': 'low'}

df.join(df.rename(columns=d)
         .groupby(level=0, axis=1).sum()
         .reindex(['low','mid','high'], axis=1, fill_value=0)
         .add_prefix('total_')
       )

Output:

   company1  company2  company3  total_low  total_mid  total_high
0         1         0         0          0          0           1
1         0         0         1          1          0           0
2         0         1         1          2          0           0

user_8986975 · Answer

Not as short as @Quang Hoang 's but Another way;

Melt dataframe

df2=pd.melt(df,  value_vars=['company1', 'company2', 'company3'])

Map dictionary creating another column total

df2['total']=df2.variable.map(d)

Pivot high, low and add middle and join to df

compa=['low','medium','high']
df.join(df2.groupby(['variable','total'])['value'].sum().unstack('total', fill_value=0).reindex(compa,axis=1, fill_value=0).add_prefix('total_').reset_index().drop(columns=['variable']))

How to match column names with dictionary keys and add value to counter

2 Answers