Python使用Matplotlib按条件设置子图背景色时代码报错如何解决？

• 核心报错原因：df.query(f'wagon == {i}')['dia'][0]的取值方式错误，query返回的Series会保留原DataFrame的行索引，比如筛选wagon=3的行时，返回的行索引是4、5，直接用[0]索引会触发KeyError，不存在对应索引的元素。
• 冗余操作：代码中重复调用query筛选同一wagon的数据集，无意义增加运算量。

import pandas as pd
import matplotlib.pyplot as plt

df = pd.DataFrame({'wagon' : [1, 1, 2, 2, 3, 3], 'ind' : [55, 55, 65, 65, 78, 78], 'axl1' : [1, 3, 1, 9, 1, 5], 
                   'axl2' : [1, 4, 1, 5, 1, 12], 'dia' : [2, 2, 1, 1, 3, 3], 'defects' : ['диаметр', 'диаметр', 'перекос', 'перекос', 'обод', 'обод']})

wags = df['wagon'].unique()
inds = df['ind'].unique()
defect = df['defects'].unique()

for (i, j, k) in tuple(zip(wags, inds, defect)):
    # 提前提取当前wagon对应的子数据集，避免重复查询
    sub_df = df[df['wagon'] == i]
    fig, axs = plt.subplots(1, 2, figsize = (10, 5)) # 调整宽度避免子图挤压

    axs[0].plot(sub_df['axl1'].values, marker = 'o')
    axs[0].set_xlabel('distance')
    
    axs[1].plot(sub_df['axl2'].values, marker = 'o')
    # 用iloc[0]取筛选结果的第一个元素，不受原索引影响
    if sub_df['dia'].iloc[0] > 2:
        axs[1].set_facecolor('#e8f4f8') # 浅蓝背景，也可按需替换为其他颜色
    
    axs[1].set_xlabel('distance')  
    fig.suptitle(f'Wagon = {i}, Defect = {k}', fontsize = 16)

    plt.show()

• 新增子数据集提前提取逻辑，减少重复的筛选操作
• 用.iloc[0]取筛选后子数据集的第一个dia值，完全避开原DataFrame索引的影响
• 调整了画布宽度避免子图内容挤压，背景色替换为更适合查看数据的浅蓝，你也可以按需改成其他有效色值
• 把重复的axs[1]绘图逻辑移到判断外，减少代码冗余

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