CODEWITHSUNDEEP
pythonpandasmatplotlibscatter-plot
user0
user0

Reputation: 195

Scatterplot using different colors for different categories

I would like to scatterplot a pandas DataFrame with columns 'x', 'y', 'type' (and others columns, and thousands of lines) where x, y are floats and 'type' is a string, which is either, say, "pear", "apple" or "banana". I want to draw a point at (x, y) with color according to the type.

I tried:

df.plot.scatter(x='x', y='y', c='type', colormap=cmap)

with a definition like

cmap = colormap_from_dict({ 'pear': 'red', 'apple': 'blue', 'banana': 'yellow' })

But I cannot find such a function "colormap_from_dict"... Do you have a solution? Thanks.

For the record, the DataFrame looks like this:

df = pd.DataFrame([[1, 2, 'banana'], [2, 1, 'apple'], [3, 3, 'pear']], columns=['x', 'y', 'type'])

Upvotes: 1

Views: 2344

Answers (1)

JohanC
JohanC

Reputation: 80279

You can apply the color via a lambda function.

Try this:

from matplotlib import pyplot as plt
import pandas as pd

df = pd.DataFrame([[1, 2, 'banana'], [2, 1, 'apple'], [3, 3, 'pear']], columns=['x', 'y', 'type'])
colors = { 'pear': 'red', 'apple': 'blue', 'banana': 'yellow' }

df.plot.scatter(x='x', y='y', c=df['type'].apply(lambda x: colors[x]))
plt.show()

Upvotes: 3

Related Questions