Classify DataFrame rows based on first matching condition

Question

I have a pandas DataFrame, each column represents a quarter, the most recent quarters are placed to the right, not all the information gets at the same time, some columns might be missing information (NaN values)

I would like to create a new column with the first criteria number that the row matches, or zero if it doesn't match any criteria

The criteria gets applied to the 3 most recent columns that have data (an integer, ignoring NaNs) and a match is considered if the value in the list is greater than or equal to its corresponding value in the DataFrame

I tried using apply, but I couldn't make it work and the failed attempts were slow

import pandas as pd
import numpy as np

criteria_dict = {
    1: [10, 0, 10]
    , 2: [0, 10, 10]
    }

list_of_tuples = [
    (78, 7, 11, 15),  # classify as 2 since  7 >= 0, 11 >= 10, 15 >= 10
    (98, -5, np.NaN, 18), # classify as 0, ignoring NaN it doesn't match any criteria because of the -5
    (-78, 20, 64, 28),    # classify as 1  20 >= 10, 64 >= 0, 28 >= 10
    (35, 63, 27, np.NaN), # classify as 1, NaN value should be ignored, 35 >= 10, 63 >=0, 27 >= 10
    (-11, 0, 56, 10) # classify as 2,   0 >= 0, 56 >= 10, 10 >= 10
]

df = pd.DataFrame(
    list_of_tuples,
    index=['A', 'B', 'C', 'D', 'E'],
    columns=['2021Q2', '2021Q3', '2021Q4', '2022Q1']
)

print(df)

Answer 1

Probably using apply is the most straightforward, but I wanted to try a solution with numpy, which should be faster with dataframes with many rows.

import numpy as np

# Rows with too many NaNs.
df_arr = df.to_numpy()

# Find NaNs.
nans = np.nonzero(np.isnan(df_arr))

# Roll the rows so that the latest three columns with valid data are all to the right.
for row, col in zip(*nans):
    df_arr[row, :] = np.roll(df_arr[row, :], shift=4-col)

# Check for matching criteria.
df['criteria'] = np.select([np.all((df_arr[:, 1:] - criteria_dict[crit])>=0, axis=1) for crit in criteria_dict],
                      [crit for crit in criteria_dict])    

print(df)

   2021Q2  2021Q3  2021Q4  2022Q1  criteria
A      78       7    11.0    15.0       2.0
B      98      -5     NaN    18.0       0.0
C     -78      20    64.0    28.0       1.0
D      35      63    27.0     NaN       1.0
E     -11       0    56.0    10.0       2.0

Some timings on df = pd.concat([df]*10000):

# 103 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit numpy(df)

# 1.32 s ± 14.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit pandas_apply(df)

So it is ~10x faster.

Answer 2

Applying a custom function to each row should work.

def func(x):
    x = x.dropna().to_numpy()[-3:]
    if len(x) < 3:
        return 0
    for k, v in criteria_dict.items():
        if np.all(x >= v):
            return k
    return 0

df.apply(func, axis=1)

Answer 3

It is possible to achieve a full vectorial comparison. Note that the bottleneck is the broadcasting step that will create an intermediate array of K*N*M size where M*N is the size of the subset of the dataframe (here 5*3) and K*N that of the criterions (here 2*3). You need to have enough memory to create this array.

Step by step procedure:

First get last 3 non-nan values as b:

N = 3

a = df.to_numpy()

from scipy.stats import rankdata
b = a[rankdata(~np.isnan(a), method='ordinal', axis=1)>(a.shape[1]-N)].reshape(-1,N)

array([[ 7., 11., 15.],
       [98., -5., 18.],
       [20., 64., 28.],
       [35., 63., 27.],
       [ 0., 56., 10.]])

Then craft an array with the conditions as c;

c = np.array(list(criteria_dict.values()))

array([[10,  0, 10],
       [ 0, 10, 10]])

Broadcast the comparison of b and c and get all values >=:

d = (b>=c[:, None]).all(2)

array([[False, False,  True,  True, False],
       [ True, False,  True,  True,  True]])

Get index of first True using the criteria_dict keys (else 0):

e = np.where(d.any(0), np.array(list(criteria_dict))[np.argmax(d, axis=0)], 0)

array([2, 0, 1, 1, 2])

Assign to DataFrame:

df['criteria'] = e 

   2021Q2  2021Q3  2021Q4  2022Q1  criteria
A      78       7    11.0    15.0         2
B      98      -5     NaN    18.0         0
C     -78      20    64.0    28.0         1
D      35      63    27.0     NaN         1
E     -11       0    56.0    10.0         2