Pandas堆叠条形图颜色映射遇缺失类别报错的解决方法

问题场景 #

给定以下DataFrame：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({
    'year': [2022,2022,2022,2022,2022,2023,2023],
    'source': ['youtube', 'youtube', 'facebook', 'facebook', 'facebook', 'google', 'google'], 
             'score': [10,20,100,200,300,90,70], 
             'rating': ['small', 'large', 'small', 'medium', 'large', 'medium', 'large']})

需求是为source列的每个值生成独立的堆叠条形图，通过color参数传入固定颜色映射：

color_map = {'small': 'yellow', 'medium': 'green', 'large':'blue'}

但部分source（如youtube）缺少medium类别，导致绘图时因颜色映射的键与数据列不匹配而报错。原代码如下：

for company in df['source'].unique(): 
    # filter to make plot of company only
    temp_df = df[df.source == company]
    
    temp_df = temp_df.pivot_table(index=temp_df.year, columns=['source', 'rating'], values='score', aggfunc='sum')
    
    color_map = {'small': 'yellow', 'medium': 'green', 'large':'blue'}
    
    fig, ax = plt.subplots(1,1)
    
    ax = temp_df.plot.bar(stacked=True, figsize=(10, 6), ylabel='scores', xlabel='dates', title='Scores', ax = ax,color = color_map)

解决方案 #

核心思路是确保每个source对应的数据集包含所有rating类别（缺失类别填充0），让颜色映射的键与数据列完全匹配，避免报错。具体修改步骤如下：

1. 提前定义所有可能的rating类别，保证统一顺序
2. 处理临时数据集时，补充缺失的rating类别并填充0
3. 简化pivot结构，去掉冗余的source列层级

修改后的完整代码：

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({
    'year': [2022,2022,2022,2022,2022,2023,2023],
    'source': ['youtube', 'youtube', 'facebook', 'facebook', 'facebook', 'google', 'google'], 
             'score': [10,20,100,200,300,90,70], 
             'rating': ['small', 'large', 'small', 'medium', 'large', 'medium', 'large']})

# 定义所有rating类别，确保和颜色映射顺序一致
all_ratings = ['small', 'medium', 'large']
color_map = {'small': 'yellow', 'medium': 'green', 'large':'blue'}

for company in df['source'].unique():
    # 过滤当前source的数据
    temp_df = df[df['source'] == company]
    
    # 按year和rating分组求和，强制保留所有rating类别，缺失值填充0
    temp_df = temp_df.pivot_table(
        index='year',
        columns='rating',
        values='score',
        aggfunc='sum',
        fill_value=0
    ).reindex(columns=all_ratings, fill_value=0)
    
    # 绘制堆叠条形图
    fig, ax = plt.subplots(figsize=(10, 6))
    temp_df.plot.bar(
        stacked=True,
        ylabel='scores',
        xlabel='year',
        title=f'Scores for {company}',
        ax=ax,
        color=[color_map[rating] for rating in all_ratings]
    )
    plt.tight_layout()
    plt.show()

关键修改说明 #

• 移除原pivot中冗余的source列层级，因为已经按source过滤，无需重复保留
• 使用reindex(columns=all_ratings)强制保留所有rating类别，缺失类别填充0，确保列与颜色映射的键完全匹配
• 按rating顺序取对应颜色，避免因列顺序不一致导致颜色错位
• 标题添加source名称，让每张图的主题更清晰

内容的提问来源于stack exchange，提问作者Beans On Toast