Vectorized way to find first column matching criteria in a Pandas DataFrame

user_5858851 0
Viewed 216

Suppose I have the following pandas DataFrame:

          A         B         C
0  0.548814  0.791725  0.978618
1  0.715189  0.528895  0.799159
2  0.602763  0.568045  0.461479
3  0.544883  0.925597  0.780529
4  0.423655  0.071036  0.118274
5  0.645894  0.087129  0.639921
6  0.437587  0.020218  0.143353
7  0.891773  0.832620  0.944669
8  0.963663  0.778157  0.521848
9  0.383442  0.870012  0.414662

Which can be created with the following code:

import pandas as pd
import numpy as np

size = 10
np.random.seed(0)
keys = ["A", "B", "C"]
df = pd.DataFrame({k: np.random.random(size) for k in keys})

How can I find the first column that meets a given criteria?

In this case, suppose my criteria is that I want the first column in which the value is less than some p, say 0.5. If no column meets this criteria, I want to return "No Match".

Using apply, this could do done as follows:

p = 0.5
first = df.apply(
    lambda row: next((x for i, x in enumerate(df.columns) if row[x]<p), "No Match"), 
    axis=1
)
print(first)
#0    No Match
#1    No Match
#2           C
#3    No Match
#4           A
#5           B
#6           A
#7    No Match
#8    No Match
#9           A
#dtype: object

Is there a more efficient (vectorized) way to do this? I am thinking there should be some way using argmax(), but I haven't gotten it to work.

Also, I'm using pandas 0.19.2 and I'm not sure if I can upgrade.

print(pd.__version__)
#u'0.19.2'
pandas python
2 Answers

You can use NumPy argmax, but will need to overwrite instances where your condition is never met in a given row:

mask = df.lt(0.5)
df['first'] = np.where(mask.any(1), df.columns[mask.values.argmax(1)], 'No Match')

You can also use Pandas idxmax:

df['first'] = np.where(mask.any(1), mask.idxmax(1), 'No Match')

print(df)

          A         B         C     first
0  0.548814  0.791725  0.978618  No Match
1  0.715189  0.528895  0.799159  No Match
2  0.602763  0.568045  0.461479         C
3  0.544883  0.925597  0.780529  No Match
4  0.423655  0.071036  0.118274         A
5  0.645894  0.087129  0.639921         B
6  0.437587  0.020218  0.143353         A
7  0.891773  0.832620  0.944669  No Match
8  0.963663  0.778157  0.521848  No Match
9  0.383442  0.870012  0.414662         A
user_92095460

IIUC with dot 

df.lt(0.5).dot(df.columns).str[0].fillna('notmatch')
Out[167]: 
0    notmatch
1    notmatch
2           C
3    notmatch
4           A
5           B
6           A
7    notmatch
8    notmatch
9           A
dtype: object
user_79645270
Related