如何基于多数据集数组绘制带线性趋势线的性别分组散点图

问题分析 #

你的代码存在两处关键错误：

• 性别筛选条件错误：数据中性别标识为'm'/'f'，但代码中用了'male'/'female'，导致无法正确分组
• 线性拟合时未定义x和y变量，np.polyfit无法执行

修正后的完整代码 #

import matplotlib.pyplot as plt
import numpy as np

# 给定数据
vals = [[39, 50907.00500830538, 'm'], [71, 58137.09607273632, 'm'], [27, 44311.25956375814, 'f'], [50, 53194.40398297405, 'f'], [41, 48227.6226667045, 'f'], [38, 51081.77610221989, 'f'], [25, 49202.743772155154, 'f'], [45, 46958.227355122865, 'm'], [46, 54815.07514726054, 'm'], [25, 46734.0863416376, 'f'], [44, 52252.36769285552, 'm'], [70, 58453.80544624214, 'f']]

fig, ax = plt.subplots()

# 修正性别筛选条件，使用'm'/'f'匹配数据
male_data = [(a,i) for a,i,g in vals if g == 'm']
male_ages = [a for a,i in male_data]
male_incomes = [i for a,i in male_data]
ax.scatter(male_ages, male_incomes, color='blue', label='male')

female_data = [(a,i) for a,i,g in vals if g == 'f']
female_ages = [a for a,i in female_data]
female_incomes = [i for a,i in female_data]
ax.scatter(female_ages, female_incomes, color='red', label='female')

# 为男性数据添加线性趋势线
if male_ages:
    z_male = np.polyfit(male_ages, male_incomes, 1)
    p_male = np.poly1d(z_male)
    ax.plot(male_ages, p_male(male_ages), color='blue', linestyle='--', label='Male Trend')

# 为女性数据添加线性趋势线
if female_ages:
    z_female = np.polyfit(female_ages, female_incomes, 1)
    p_female = np.poly1d(z_female)
    ax.plot(female_ages, p_female(female_ages), color='red', linestyle='--', label='Female Trend')

ax.legend()
ax.set_xlabel('age')
ax.set_ylabel('income')

plt.show()

修改说明 #

1. 修正性别筛选：将g == 'male'改为g == 'm'，g == 'female'改为g == 'f'，确保正确匹配数据中的性别标识
2. 定义拟合变量：分别使用male_ages和male_incomes、female_ages和female_incomes作为np.polyfit的输入变量，避免未定义错误
3. 添加分组趋势线：为男女数据分别计算并绘制线性趋势线，用虚线区分散点和趋势线，同时添加对应的图例标签
4. 空数据判断：加入if male_ages:和if female_ages:的判断，避免某一性别无数据时出现拟合错误

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