如何改写Plotly Treemap的color_discrete_map适配动态类别数量

我正在开发一个仪表盘，使用如下数据集构建Treemap：

import pandas as pd
import plotly.express as px

col1 = ['United States','Mexico','Canada','Argentina','France','Portugal','Italy','Spain','Brazil','Peru']
col2 = [4,6,6,9,10,3,20,15,16,2]

df = pd.DataFrame({
    'country': col1,
    'count': col2
})

我将使用过滤器来改变显示的国家及对应数值，以下是生成Treemap的代码：

tree_fig = px.treemap(
    df, 
    path = ['country'],
    values = 'count',
    template='plotly_dark',
    color = 'country',
    color_discrete_map={
        df['country'][0]: '#626ffb', 
        df['country'][1]: '#b064fc', 
        df['country'][2]: '#ef563b',
        df['country'][3]: '#f45498',
        df['country'][4]: '#ff94fc',
        df['country'][5]: '#a8f064',
        df['country'][6]: '#24cce6',
        df['country'][7]: '#ffa45c',
        df['country'][8]: '#00cc96',
        df['country'][9]: '#fff000'

    }   
)
tree_fig.update_traces(
    hovertemplate='Count=%{value}'
)
tree_fig.show()

当前代码可正常运行，但使用过滤器后，DataFrame中的国家数量可能少于10个，请问如何改写color_discrete_map参数以适配类别数量的变化？

核心思路是先定义全局固定的国家-颜色映射字典，再根据过滤后的DataFrame提取对应国家的颜色键值对，自动适配类别数量变化。

1. 定义全局国家-颜色映射 #

提前把所有可能涉及的国家和对应颜色定义好，避免重复编写：

# 全局固定映射，包含所有可能出现的国家
country_color_map = {
    'United States': '#626ffb',
    'Mexico': '#b064fc',
    'Canada': '#ef563b',
    'Argentina': '#f45498',
    'France': '#ff94fc',
    'Portugal': '#a8f064',
    'Italy': '#24cce6',
    'Spain': '#ffa45c',
    'Brazil': '#00cc96',
    'Peru': '#fff000'
}

2. 生成适配过滤后DataFrame的颜色映射 #

从全局映射中筛选出当前DataFrame里存在的国家对应的颜色：

# 示例：假设这是经过过滤器处理后的DataFrame
filtered_df = df[df['count'] > 5]  # 可替换为你的实际过滤逻辑

# 生成仅包含当前显示国家的颜色映射
filtered_color_map = {country: country_color_map[country] for country in filtered_df['country'].unique()}

3. 修改Treemap生成代码 #

将原来的color_discrete_map替换为适配后的映射：

tree_fig = px.treemap(
    filtered_df, 
    path=['country'],
    values='count',
    template='plotly_dark',
    color='country',
    color_discrete_map=filtered_color_map  # 使用适配后的映射
)
tree_fig.update_traces(
    hovertemplate='Count=%{value}'
)
tree_fig.show()

可选优化：兼容未定义颜色的国家 #

如果后续可能新增未在全局映射中定义的国家，可以添加默认颜色避免报错：

filtered_color_map = {
    country: country_color_map.get(country, '#888888')  # 未定义的国家默认用灰色
    for country in filtered_df['country'].unique()
}

这样无论过滤器如何调整，color_discrete_map都会自动匹配当前显示的国家，不会出现键不匹配或多余键的问题。

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