如何为带年月索引的Pandas DataFrame按采样年份分配对应颜色？

问题重现 #

你尝试为DataFrame中的每个年份分配唯一XKCD颜色，但当前代码执行后生成了多列<lambda>结果，无法得到每个年份对应单一颜色的预期输出。

问题原因 #

你的代码中，apply方法里误用了生成器表达式而非单个lambda函数：

df['year_color'] = df['Year'].apply(lambda x: clr if x==year else np.nan for year,clr in zip(years_list,req_colors_list))

这个生成器会遍历所有年份-颜色对，每一对生成一个lambda结果，最终被Pandas解析为多列数据，而非单一列的颜色值。

解决方案 #

核心思路是先创建年份到颜色的映射字典，再通过字典查找为每个年份匹配对应颜色，既高效又能得到单一列的结果。

修正后的完整代码

import matplotlib.colors as mcolors
import pandas as pd
import numpy as np

# 定义需要的XKCD颜色列表
colors_list = [
    'xkcd:cloudy blue',
    'xkcd:dark pastel green',
    'xkcd:dust',
    'xkcd:electric lime',
    'xkcd:fresh green',
    'xkcd:light eggplant'
    # 可继续添加更多颜色
]

# 示例DataFrame
df = pd.DataFrame({
    'sensor_value': [5171.318942, 5085.094086, 5685.681944, 6097.877688, 6063.909946],
    'Year': [2002, 2002, 2004, 2006, 2003],
    'Month': [4,5,6,7,8]
})

# 获取唯一年份列表并匹配对应数量的颜色
years_list = df['Year'].unique().tolist()
req_colors_list = colors_list[:len(years_list)]

# 创建年份-颜色映射字典
year_color_map = dict(zip(years_list, req_colors_list))

# 为每个年份分配颜色（生成单一列）
df['year_color'] = df['Year'].map(year_color_map)

# 查看结果
print(df[['Year', 'year_color']])

预期输出

Year         year_color
0  2002    xkcd:cloudy blue
1  2002    xkcd:cloudy blue
2  2004  xkcd:dark pastel green
3  2006      xkcd:fresh green
4  2003            xkcd:dust

补充说明 #

• 使用dict(zip(...))创建映射字典是最直观且高效的方式，字典查找的时间复杂度为O(1)，远优于循环匹配。
• 如果需要处理未在映射字典中的年份，可以用map结合get方法设置默认值，例如：df['Year'].map(lambda x: year_color_map.get(x, 'xkcd:white'))。

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