-1

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.

Ajax
  • 159
  • 7

1 Answers1

0

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)
Ben.T
  • 29,160
  • 6
  • 32
  • 54
  • and what if there is a condition. I have value in original column from 0 to 100 and i want that 1 and 2 stay unchanged but all values which are 3 or greater than 3 convert to 3 – Ajax Apr 09 '20 at 01:53
  • and get the error "method object is not subscriptable" – Ajax Apr 09 '20 at 01:57
  • @Arpit I don't get the error, so you need to elaborate if you want help. and for your previous comment, then it could be `np.where(df[col]<3, df[col], 3)` – Ben.T Apr 09 '20 at 02:08