Create dataframe column with elif of different types in pandas

Question

This is an extension of question.

I'd like to do some if/elif/else logic to create a dataframe column, pseudocode example:

if Col1 = 'A' and Col2 = 1 then Col3 = 'A1'
else if Col1 = 'A' and Col2 = 0 then Col3 = 'A0'
else Col3 = 'XX'

is it ok to mix up the types there? I'm getting this error:

TypeError: cannot compare a dtyped [int64] array with a scalar of type [bool]

You use single `=` sign to compare instead of `==` I think that my cause an issue — Tom Ron, Aug 31 '16 at 10:55
yes, sorry, was just using pseudocode. should have done if df.['Col1'] == 'A' — Angus, Aug 31 '16 at 12:45

jezrael · Accepted Answer · 2016-08-31T11:01:03.663

I think you can use:

df['Col3'] = 'XX'
df.loc[(df.Col1 == 'A') & (df.Col2 == 1), 'Col3'] = 'A1'
df.loc[(df.Col1 == 'A') & (df.Col2 == 0), 'Col3'] = 'A0'

With double numpy.where:

df["Col3"] = np.where((df.Col1 == 'A') & (df.Col2 == 1) , "A1", 
             np.where((df.Col1 == 'A') & (df.Col2 == 0), 'A0', 'XX'))

Sample:

df = pd.DataFrame({'Col1':['A','B','A','B'],
                   'Col2':[1,1,0,0]})

print (df)
  Col1  Col2
0    A     1
1    B     1
2    A     0
3    B     0

df['Col3'] = 'XX'
df.loc[(df.Col1 == 'A') & (df.Col2 == 1), 'Col3'] = 'A1'
df.loc[(df.Col1 == 'A') & (df.Col2 == 0), 'Col3'] = 'A0'

df["Col4"] = np.where((df.Col1 == 'A') & (df.Col2 == 1) , "A1", 
             np.where((df.Col1 == 'A') & (df.Col2 == 0), 'A0', 'XX'))

print (df)
  Col1  Col2 Col3 Col4
0    A     1   A1   A1
1    B     1   XX   XX
2    A     0   A0   A0
3    B     0   XX   XX

Create dataframe column with elif of different types in pandas

1 Answers1