python - How to change background cell color according with a given key in pandas?

Question

I have the following dataframe:

df = pd.DataFrame({"A":["foo", "foo", "foo", "bar"],  
                   "B":["A","A","B","A"], 
                   "C":[0,3,1,1]})

How can I change the cell color of columns A and B, grouped by its values. I mean, this would be the desirable output:

Maybe something like:

df.groupby(by=['A', 'B']).style.change_background()

As the real dataframe has hundreds of rows, I'd be interestested in assigning the colors automatically.

Answer 1

from matplotlib.pyplot import cm 
from matplotlib.colors import to_hex

# convert groups to indices    
g = pd.Categorical(df.A + df.B).codes

# generate a list hex colors
color = cm.rainbow(g / g.max())
hex_color = [to_hex(col) for col in color]
print('generate colors: ', hex_color)

# apply the colors to style for columns A and B
df.style.apply(
    lambda col: ['background-color: %s' % (hex_color[i]) for i in range(col.size)], 
    subset=['A', 'B']
)

Update by Antonvbr: (Insipred by this solution I added a solution using seaborn too).

import seaborn as sns
import pandas as pd

df = pd.DataFrame({"A":["foo", "foo", "foo", "bar"],  
                   "B":["A","A","B","A"], 
                   "C":[0,3,1,1]})

g = pd.Categorical(df.A + df.B).codes # convert groups to indices  
n = np.unique(g).size
palette = sns.color_palette("hls", n).as_hex()

# apply the colors to style for columns A and B
df.style.apply(
    lambda x: ['background-color: {}'.format(palette[i]) for i in g], 
    subset=['A', 'B']
)