Optimize Apply to create an list of string with Boolean columns

user_12316701 0
Viewed 42

I must check to a boolean value in rows, get the name of the columns to build a list and create a new column with this list of strings.

I have this code, it works perfectly but it's very slow (I used 125 000 rows for 20 columns). Do you have an idea to optimize this code?

import pandas as pd 

def exlusions(x: pd.Series) -> pd.Series:
        reasons = [name for name in x.index if x[name] == True]
        return ",".join(reasons)

dct_data = {
    'A' : [False, False, True, False, True, False],
    'B' : [False, False, False,False, False, False],
    'C' : [False, True, False,False, False, False],
    'D' : [False, False, True, False, False, False],
    'Client' : ['Paul', 'Nick', 'Josh', 'Flo', 'Julia', 'Lucia']
}

df = pd.DataFrame(dct_data)
df = df[list(df.select_dtypes(include='bool').columns) + ['Client']]
df = df[df[list(df.select_dtypes(include='bool').columns)].any(1)]
df['Exclusions'] = df.apply(lambda x: exlusions(x), axis=1)

df
pandas apply lambda python
1 Answers

We can use DataFrame.dot with the columns and str.rstrip the extra comma:

# Boolean Columns
cols = df.columns[df.dtypes == 'bool']
# Filter DataFrame rows
df = df[df[cols].any(axis=1)]
# Take dot product (multiply then sum rows) and remove trailing comma
df['Exclusions'] = df[cols].dot(cols + ',').str.rstrip(',')

df:

       A      B      C      D Client Exclusions
1  False  False   True  False   Nick          C
2   True  False  False   True   Josh        A,D
4   True  False  False  False  Julia          A
user_154978880
Related