pandas数据框绘制散点图时如何添加Matplotlib图例

Question

我将为您翻译以下内容：

我正在尝试在下面的图形中正确获取图例。它应该只有“绿色”，“蓝色”和“红色”，并带有相应的颜色。但是显示的东西很混乱。

代码如下：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({
        'category':['blue','green','red','blue','green','red','blue','green','red'],
        'attempts':[8955,7881,6723,100,200,300,4567,876,54],
        'success':[3000,7500,2000, 256,4567,4567,7665,543,43]
})
fig,ax = plt.subplots()
plt.scatter(df['attempts'],df['success'],c=df['category'],label=df['category'])
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

我该如何正确地实现这个？（这里有一个类似的例子：https://pythonspot.com/matplotlib-scatterplot/，第二部分"带有分组的散点图"，但它不是基于Pandas数据框架）。

- Alhpa Delta

Answer 1

您可以使用seaborn的scatterplot:

fig,ax = plt.subplots()
sns.scatterplot(data=df, hue='category', x='attempts', y='success')
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

输出:

或者使用纯matplotlib：

fig,ax = plt.subplots()
for k,d in df.groupby('category'):
    ax.scatter(d['attempts'], d['success'], label=k)
plt.legend(loc=2)
plt.savefig('scatter.png')
plt.show()

输出：

- Quang Hoang

---

我从未知道可以像这样使用groupby。您能否提供一个网址或一些文档，让我可以阅读有关此内容的信息？不仅仅是在绘图的上下文中，而是在任何其他可能有用的上下文中。 - Alhpa Delta

这是来自pandas文档，groupby()返回一个GroupBy对象，该对象具有一个__iter__()方法，允许使用for ... in。 - Quang Hoang

Answer 2

如果你想在matplotlib中使用单个散点图，代码如下所示：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

df = pd.DataFrame({
        'category':['blue','green','red','blue','green','red','blue','green','red'],
        'attempts':[8955,7881,6723,100,200,300,4567,876,54],
        'success':[3000,7500,2000, 256,4567,4567,7665,543,43]
})

u, inv = np.unique(df.category.values, return_inverse=True)
cmap = ListedColormap(u)

fig,ax = plt.subplots()
scatter = plt.scatter(df['attempts'],df['success'],c=inv, cmap=cmap)
plt.legend(scatter.legend_elements()[0], u, loc=2)
plt.savefig('scatter.png')
plt.show()

- ImportanceOfBeingErnest