How to apply a function to get the encoded specific columns in pandas

Question

I have this function:

get_class(cols):
    if cols == 1:
        return 1
    elif cols ==2:
        return 2
    else:
        return 0

I made a list of certain columns like this:

cols = ['night', 'day']
cols_en = []

for each in cols:
    each = cols + '_en'
    cols_en.append(each)

Here, i want the function get_class to apply to the cols and take the output in cols_en. I want to automate this code:

df ['night_en'] = [1 if x==1 else 2 if x==2 else 0 for x in df['night']]

Idea is to apply the function to all the columns which are in list cols and get the output, where the columns have the fuction get_class applied and output columns have _en at the end. Maybe using map function also. Any idea to achieve this? I have read several similar articles but didn't help much.

Answer 1

IIUC, you want to create a new column with 0 if the value is not 1 nor 2 in the original, for example, you can use np.where. and create a loop over your list cols to create the new column

for col in cols:
    df[f'{col}_en'] = np.where(df[col].isin([1,2]), df[col], 0)