Looking for a better solution to create conditional columns in pandas

Question

Consider the following

df:

id,flag,amt 

10,1,100 
11,0,100
13,1,100
14,0,100

current code :

def func(row):
    if row['flag'] == 1:
        val = row['amt'] * 2
    else:
        val = row['amt']
    return val

    df['op'] = df.apply(func,axis=1)

output:

df:

    id,flag,amt,op 

    10,1,100,200
    11,0,100,100
    13,1,100,200
    14,0,100,100

Is there a better way to implement this? This solution is taking time!

Answer 1

You can do this in bulk with:

df['op'] = (df['flag']+1) * df['amt']

This gives us:

>>> (df['flag']+1) * df['amt']
0    200
1    100
2    200
3    100
dtype: int64

In case df['flag'] can have other values than 0 and 1, we can use np.where(..), like:

df['op'] = np.where(df['flag'] == 1, 2*df['amt'], df['amt'])