如何为基于聚合列的州级散点图添加对应州标签？

问题背景 #

现有如下pandas DataFrame数据：

import pandas as pd
df = pd.DataFrame({"State": ["CA", "NY", "CA", "NY", "CA", "NY", "TX", "TX", "TX"],
                   "Company": ["A", "A", "A", "B", "C", "D", "A", "B", "B"],
                   "Profits": [3, 2, 5, 6, 7, 2, 2, 4, 7]})

需求是绘制散点图：每个点对应一个州，X轴为该州的唯一公司数量，Y轴为该州的平均利润。尝试了以下代码：

n_companies = df.groupby("State")["Company"].nunique()
mean_profits = df.groupby("State")["Profits"].mean()
import matplotlib.pyplot as plt
plt.scatter(n_companies, mean_profits)
plt.show()

散点图可正常生成，但需要为每个点添加对应的州标签。

解决方法 #

可以通过以下步骤实现标签添加：

• 合并分组后的指标数据，保留州名称信息
• 遍历每个州的数据，在散点位置添加文本标签

完整代码如下：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({"State": ["CA", "NY", "CA", "NY", "CA", "NY", "TX", "TX", "TX"],
                   "Company": ["A", "A", "A", "B", "C", "D", "A", "B", "B"],
                   "Profits": [3, 2, 5, 6, 7, 2, 2, 4, 7]})

# 一次性分组计算两个指标，生成包含州名的DataFrame
state_stats = df.groupby("State").agg(
    unique_companies=("Company", "nunique"),
    avg_profit=("Profits", "mean")
).reset_index()

# 绘制散点图
plt.scatter(state_stats["unique_companies"], state_stats["avg_profit"])

# 为每个散点添加州标签
for idx, row in state_stats.iterrows():
    # 偏移X坐标避免标签与散点重叠
    plt.text(row["unique_companies"] + 0.05, row["avg_profit"], row["State"])

# 添加坐标轴说明
plt.xlabel("唯一公司数量")
plt.ylabel("平均利润")
plt.show()

代码说明 #

• 使用groupby().agg()可以高效完成多指标计算，同时直接保留州名称，无需单独合并数据
• plt.text()的坐标参数可根据需求调整偏移量，避免标签遮挡散点；也可额外设置字体大小、颜色等样式

内容的提问来源于stack exchange，提问作者Smithey