Pandas create a color for each two value pair

May 26, 2022

I have a data frame with XY values which I would like to plot in a scatter plot. I would like to create groups which have a same XY values and apply the same color in the plot:

df = pd.DataFrame({'X': [1.1, 1.1, 5.5, 5.5, 5.5,10.1],
                   'Y': [3.7, 3.7, 5.6, 5.6, 5.6, 3.2],
                   'ID': ['ID1', 'ID2','ID3','ID4','ID5', 'ID6'})


X        Y         ID         color (hex or RGB)
1.1      3.7       ID1        0,255,255
1.1      3.7       ID2        0,255,255
5.5      5.6       ID3        255,0,255
5.5      5.6       ID4        255,0,255
5.5      5.6       ID5        255,0,255
10.1     3.2       ID6        0,0,255

How to map a color for each XY group?

>Solution :

You can define a list of colors. After groupby the X and Y column, get group id with ngroup() then map it to the list of colors

colors = ['0,255,255', '255,0,255', '0,0,255']

df['color'] = df.groupby(['X', 'Y']).ngroup().map(dict(enumerate(colors)))

print(df)

      X    Y   ID      color
0   1.1  3.7  ID1  0,255,255
1   1.1  3.7  ID2  0,255,255
2   5.5  5.6  ID3  255,0,255
3   5.5  5.6  ID4  255,0,255
4   5.5  5.6  ID5  255,0,255
5  10.1  3.2  ID6    0,0,255