Pandas conditional creation of a series/dataframe column
List comprehension is another way to create another column conditionally. If you are working with object dtypes in columns, like in your example, list comprehensions typically outperform most other methods.
Example list comprehension:
df['color'] = ['red' if x == 'Z' else 'green' for x in df['Set']]
%timeit tests:
import pandas as pd
import numpy as np
df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
%timeit df['color'] = ['red' if x == 'Z' else 'green' for x in df['Set']]
%timeit df['color'] = np.where(df['Set']=='Z', 'green', 'red')
%timeit df['color'] = df.Set.map( lambda x: 'red' if x == 'Z' else 'green')
1000 loops, best of 3: 239 µs per loop
1000 loops, best of 3: 523 µs per loop
1000 loops, best of 3: 263 µs per loop
If you only have two choices to select from:
df['color'] = np.where(df['Set']=='Z', 'green', 'red')
For example,
import pandas as pd
import numpy as np
df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
df['color'] = np.where(df['Set']=='Z', 'green', 'red')
print(df)
yields
Set Type color
0 Z A green
1 Z B green
2 X B red
3 Y C red
If you have more than two conditions then use np.select
. For example, if you want color
to be
yellow
when(df['Set'] == 'Z') & (df['Type'] == 'A')
- otherwise
blue
when(df['Set'] == 'Z') & (df['Type'] == 'B')
- otherwise
purple
when(df['Type'] == 'B')
- otherwise
black
,
then use
df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
conditions = [
(df['Set'] == 'Z') & (df['Type'] == 'A'),
(df['Set'] == 'Z') & (df['Type'] == 'B'),
(df['Type'] == 'B')]
choices = ['yellow', 'blue', 'purple']
df['color'] = np.select(conditions, choices, default='black')
print(df)
which yields
Set Type color
0 Z A yellow
1 Z B blue
2 X B purple
3 Y C black
Another way in which this could be achieved is
df['color'] = df.Set.map( lambda x: 'red' if x == 'Z' else 'green')
The following is slower than the approaches timed here, but we can compute the extra column based on the contents of more than one column, and more than two values can be computed for the extra column.
Simple example using just the "Set" column:
def set_color(row):
if row["Set"] == "Z":
return "red"
else:
return "green"
df = df.assign(color=df.apply(set_color, axis=1))
print(df)
Set Type color
0 Z A red
1 Z B red
2 X B green
3 Y C green
Example with more colours and more columns taken into account:
def set_color(row):
if row["Set"] == "Z":
return "red"
elif row["Type"] == "C":
return "blue"
else:
return "green"
df = df.assign(color=df.apply(set_color, axis=1))
print(df)
Set Type color
0 Z A red
1 Z B red
2 X B green
3 Y C blue
Edit (21/06/2019): Using plydata
It is also possible to use plydata to do this kind of things (this seems even slower than using assign
and apply
, though).
from plydata import define, if_else
Simple if_else
:
df = define(df, color=if_else('Set=="Z"', '"red"', '"green"'))
print(df)
Set Type color
0 Z A red
1 Z B red
2 X B green
3 Y C green
Nested if_else
:
df = define(df, color=if_else(
'Set=="Z"',
'"red"',
if_else('Type=="C"', '"green"', '"blue"')))
print(df)
Set Type color
0 Z A red
1 Z B red
2 X B blue
3 Y C green