Pandas create a color for each two value pair-CodePudding

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?

CodePudding user response：

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

CodePudding user response：

Following up on Ynjxsjmh's answer, once the column is created you can plot them using seaborn's hue parameter:

sct = sns.scatterplot(x=df['X'],y=df['Y'],hue=df['color'],c=np.array([[int(w) for w in y] for y in [x.split(",") for x in colors]])/255)
plt.legend(title="My Groups",labels=['Group 1','Group 2','Group 3'])
plt.title("My Scatterplot") 
plt.show(sct)

Outputs: