More idiomatic "if-else" replacement in pandas

user_7340317 0
Viewed 69

I've seemingly simple problem, based on condition e.g. that value in dataframe is smaller than two, change value to 1, in opposite case to 0. Kind of "if-else".

Toy exmample, input:

   a   b
0  1  -5
1  2   0
2  3  10

Output:

   a  b
0  1  1
1  0  1
2  0  0

Here is my solution:

import pandas as pd
import numpy as np


df = pd.DataFrame({'a': [1,2,3], 'b': [-5, 0, 10]})

arr = np.where(df < 2, 1, 0)
df_fin = pd.DataFrame(data=arr, index=df.index, columns=df.columns)

I don't like direct dependency on numpy and it also a little looks verbose to me. Could it be done in more cleaner, idiomatic way?

pandas numpy python
2 Answers

General solutions:

Pandas is built in numpy, so in my opinion only need import. Here is possible set values in df[:]

import numpy as np

df[:] = np.where(df < 2, 1, 0)
print (df)
   a  b
0  1  1
1  0  1
2  0  0

A bit overcomplicated if use only pandas functions:

m = df < 2
df = df.mask(m, 1).where(m, 0)

Replace to 0,1 solution:

Convert mask for map True to 1 and False to 0 by DataFrame.view or like in another answer:

df = (df < 2).view('i1')
user_29010020

Pandas' replace might be handy here:

 df.lt(2).replace({False : 0, True: 1})
Out[7]:
   a  b
0  1  1
1  0  1
2  0  0

or you just convert the booleans to integers:

df.lt(2).astype(int)
Out[9]:
   a  b
0  1  1
1  0  1
2  0  0
user_71757130
Related