Search code examples
pythonpandasdataframestyler

Pandas dataframe background color based on condition


I have a simple DataFrame consisting of three columns A, B and C. The columns can contain different values (text) but sometimes contain the word 'ALARM'. Based on how many times 'ALARM' appears in one row, I would like to color the entire row. Problem: It always colors the entire dataframe in the same color.

Here the code to reproduce the dataframe:

data = {'A': ['ALARM', 'ALARM', 'krish', 'Peter'],
    'B': ['Tom', 'ALARM', 'krish', 'ALARM'],
    'C': ['Jack', 'ALARM', 'krish', 'ALARM'],}

df = pd.DataFrame(data)

And here the code I used to color the dataframe:

def num_alarms(series):    

    red = 'background-color: red;'
    orange = 'background-color: orange;'
    yellow = 'background-color: yellow;'
    green = 'background-color: green;'

    return [red if num == 3 else orange if num == 2 else yellow if num == 1 else green for e 
    in series]


 for index, row in df.iterrows():
    num = (df == 'ALARM').sum(axis = 1)[index]   

    s = df.style.apply(num_alarms, axis = 'index', subset = slice(index))

 display(s)

NOTE: I've been working in a Jupyter notebook.


Solution

  • Use Series.map for get mapped values by dictionary and remove loop, use only Styler.apply with axis=1:

    def styled_alarms(x):    
    
        red = 'background-color: red;'
        orange = 'background-color: orange;'
        yellow = 'background-color: yellow;'
        green = 'background-color: green;'
        
        d = {0:green,1:yellow, 2:orange, 3:red}
        return np.repeat(d.get((x == 'ALARM').sum()), len(x))
     
    df.style.apply(styled_alarms, axis=1)