Vectorize operation to create multiple new columns based on different conditions in a dataframe in Python

Viewed 305

I have a dataframe which looks like below:

+-----+---------+--------+-------------+
| Id  | Record  | Type   |   Status    |
+-----+---------+--------+-------------+
| 111 |  R1     | Email  |   Open      |
| 123 |  R2     | Event  |   Accepted  |
| 234 |  R3     | Story  |   Pending   |
| 135 |  R3     | Search |   Click     |
+-----+---------+--------+-------------+

It has around 1 million rows

The Record column can have only three values i.e. 'R1', 'R2' and 'R3'.

For record R1, type is always 'Email' and status has multiple values out of which I am only interested in 'Open' and 'Click'.

For record R2, type is always 'Event' and status has multiple values out of which I am only interested in 'Open' and 'Registered'.

However for record R3, there can be multiple values of type i.e. 'search', 'story', etc. and each type has different status values that I am interested in.

I want to create new count columns on following conditions:

  • If record is 'R1', then email_cnt = 1, else 0

  • If record is 'R2', then event_cnt = 1, else 0

  • If record is 'R3' and type is 'Story', then story_cnt = 1, else 0

  • If record is 'R3' and type is 'Search', then search_cnt = 1, else 0 and so on for all different R3 types.

I have used iterrows to iterate over each row of dataframe and substitute the values. I understand that this isn't the most efficient method, but I am not able to formulate how to use apply method or do vectorization to speed up the calculations.

Please suggest if you know of a faster/better approach.

Below is my code:

df = pd.read_csv('file.csv')

# Create new metrics columns
tempcols = ['email_cnt', 'event_cnt', 'dm_cnt', 'enc_cnt', 'exp_cnt', 'orgsearch_cnt', 'orgsocial_cnt', 'paidsm_cnt', 'paidsearch_cnt', 'pd_cnt', 'smrtroom_cnt', 'stry_cnt', 'tm_cnt']

# Append new metrics in the existing campaigns dataframe
df = pd.concat([df, pd.DataFrame(columns=tempcols)])

# Values of Status column that we need
status_vals = ['Accepted', 'Call Completed', 'Commented', 'Declined', 'Liked', 'Responded', 'Shared']

for index, row in df.iterrows():
    if((row['Record']=='R1') & (row['Status'] in ['Open', 'Click'])):
        df.at[index, 'email_cnt'] = 1
    
    if((row['Record']=='R2') & (row['Status'] in ['Open', 'Registered'])):
        df.at[index, 'event_cnt'] = 1
    
    if(row['Record']=='R3'):
        if((row['Type']=='Story') & (row['Status'] in status_vals)):
            df.at[index, 'stry_cnt'] = 1
        if((row['Type']=='Search') & (row['Status'] in status_vals+['Downloaded', 'Registered'])):
            df.at[index, 'search_cnt'] = 1
        if((row['Type']=='Experience') & (row['Status'] in status_vals)):
            df.at[index, 'exp_cnt'] = 1

df.fillna(0, inplace=True)
2 Answers

You can make a list of your conditions and make a DataFrame out of it:

conditions=[
    (df['Record'].eq('R1') & (df['Status'].isin(['Open', 'Click']))),
    (df['Record'].eq('R2') & (df['Status'].isin(['Open', 'Registered']))),
    (df['Record'].eq('R3') & df['Type'].eq('Story')),
    (df['Record'].eq('R3') & df['Type'].eq('Search'))
           ]

out=pd.concat(conditions,axis=1).astype(int)
out.columns=['email_cnt', 'event_cnt', 'story_cnt', 'search_cnt']

OR

you can use np.where() 4 times:

import numpy as np

df['email_cnt']=np.where((df['Record'].eq('R1') & (df['Status'].isin(['Open', 'Click']))),1,0)
df['event_cnt']=np.where((df['Record'].eq('R2') & (df['Status'].isin(['Open', 'Registered']))),1,0)
df['story_cnt']=np.where((df['Record'].eq('R3') & df['Type'].eq('Story')),1,0)
df['search_cnt']=np.where( (df['Record'].eq('R3') & df['Type'].eq('Search')),1,0)

Note: you can also use loc like np.where in 4 steps

You can create all the new columns with:

new_cols = [col.lower() + '_cnt' for col in set(df['Type'])]
df = df.assign({col: 0 for col in new_cols})

And then fill the 1's with your specific conditions with commands like:

df.loc[(df['Record'] == 'R1') & (df['Type'] == 'Story') & df['Status'].isin(['Open', 'Click']), 'email_cnt'] = 1

And so for all the conditions you have.

Related