How to create categorical variable based on a numerical variable
2 ways, use a couple loc
calls to mask the rows where the conditions are met:
In [309]:
df.loc[(df['col1'] > 0) & (df['col1']<= 10), 'col2'] = 'xxx'
df.loc[(df['col1'] > 10) & (df['col1']<= 50), 'col2'] = 'yyy'
df.loc[df['col1'] > 50, 'col2'] = 'zzz'
df
Out[309]:
col1 col2
0 1 xxx
1 1 xxx
2 4 xxx
3 5 xxx
4 6 xxx
5 6 xxx
6 30 yyy
7 20 yyy
8 80 zzz
9 90 zzz
Or use a nested np.where
:
In [310]:
df['col2'] = np.where((df['col1'] > 0) & (df['col1']<= 10), 'xxx', np.where((df['col1'] > 10) & (df['col1']<= 50), 'yyy', 'zzz'))
df
Out[310]:
col1 col2
0 1 xxx
1 1 xxx
2 4 xxx
3 5 xxx
4 6 xxx
5 6 xxx
6 30 yyy
7 20 yyy
8 80 zzz
9 90 zzz
You could use pd.cut
as follows:
df['col2'] = pd.cut(df['col1'], bins=[0, 10, 50, float('Inf')], labels=['xxx', 'yyy', 'zzz'])
Output:
col1 col2
0 1 xxx
1 1 xxx
2 4 xxx
3 5 xxx
4 6 xxx
5 6 xxx
6 30 yyy
7 20 yyy
8 80 zzz
9 90 zzz
You could first create a new column col2
, and update its values based on the conditions:
df['col2'] = 'zzz'
df.loc[(df['col1'] > 0) & (df['col1'] <= 10), 'col2'] = 'xxx'
df.loc[(df['col1'] > 10) & (df['col1'] <= 50), 'col2'] = 'yyy'
print df
Output:
col1 col2
0 1 xxx
1 1 xxx
2 4 xxx
3 5 xxx
4 6 xxx
5 6 xxx
6 30 yyy
7 20 yyy
8 80 zzz
9 90 zzz
Alternatively, you can also apply a function based on the column col1
:
def func(x):
if 0 < x <= 10:
return 'xxx'
elif 10 < x <= 50:
return 'yyy'
return 'zzz'
df['col2'] = df['col1'].apply(func)
and this will result in the same output.
The apply
approach should be preferred in this case as it is much faster:
%timeit run() # packaged to run the first approach
# 100 loops, best of 3: 3.28 ms per loop
%timeit df['col2'] = df['col1'].apply(func)
# 10000 loops, best of 3: 187 µs per loop
However, when the size of the DataFrame is large, the built-in vectorized operations (i.e. with the masking approach) might be faster.