如何在Pandas Styler中正确子集化多索引列以应用样式？

问题场景 #

有如下带多索引列的DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(
    {
        ("A", "Current"): [1, 2, 3],
        ("A", "Prev"): [1, 1, 1],
        ("B", "Current"): [7, 8, 9],
        ("B", "Prev"): [5, 4, 3],
    }
)

需要通过自定义的movement函数，结合对应Prev列的值来格式化Current列，且要批量处理多组类似的列对（比如A、B、C等的Current/Prev组合）。自定义函数代码如下：

def movement(row, current_col, prev_col, color_one, color_two):
    colors = pd.DataFrame("", index=row.index, columns=row.columns)
    colors[current_col] = np.where(
        (row[current_col] > 0) & (row[current_col] <= 2) & (row[prev_col] > 0),
        f"color:{color_one}",
        f"color:{color_two}",
    )
    return colors

但调用style.apply()并传入subset参数时，触发了KeyError：

df_style = df.style.apply(
    movement,
    current_col=("A", "Current"),
    prev_col=("A", "Prev"),
    color_one="green",
    color_two="red",
    axis=None,
    subset=[("A", "Current")],
)

df_style.render()

错误原因 #

当你指定subset=[("A", "Current")]时，style.apply()只会把仅包含该列的子集DataFrame传入movement函数，但你的函数里仍尝试访问原DataFrame的("A", "Prev")列——这列不在子集范围内，因此直接抛出KeyError。

解决方案 #

方法1：修改函数，从原DataFrame获取对比值 #

既然subset会截断传入的数据，我们可以让函数直接从原DataFrame中读取Prev列的值，避免依赖子集数据：

def movement(row, current_col, prev_col, color_one, color_two, original_df):
    # 从原DataFrame提取对应行的Prev列值
    prev_vals = original_df.loc[row.index, prev_col]
    colors = pd.DataFrame("", index=row.index, columns=row.columns)
    colors[current_col] = np.where(
        (row[current_col] > 0) & (row[current_col] <= 2) & (prev_vals > 0),
        f"color:{color_one}",
        f"color:{color_two}",
    )
    return colors

# 应用样式时传入原DataFrame
df_style = df.style.apply(
    movement,
    current_col=("A", "Current"),
    prev_col=("A", "Prev"),
    color_one="green",
    color_two="red",
    original_df=df,
    axis=None,
    subset=[("A", "Current")],
)

方法2：去掉subset，在函数中控制仅修改目标列 #

如果不想额外传入原DataFrame，也可以直接去掉subset参数，让函数返回完整的样式DataFrame，但仅对目标Current列设置格式，其他列留空：

def movement(row, current_col, prev_col, color_one, color_two):
    colors = pd.DataFrame("", index=row.index, columns=row.columns)
    colors[current_col] = np.where(
        (row[current_col] > 0) & (row[current_col] <= 2) & (row[prev_col] > 0),
        f"color:{color_one}",
        f"color:{color_two}",
    )
    return colors

# 直接应用样式，无需subset
df_style = df.style.apply(
    movement,
    current_col=("A", "Current"),
    prev_col=("A", "Prev"),
    color_one="green",
    color_two="red",
    axis=None,
)

方法3：循环迭代处理多组列（适配实际批量场景） #

针对实际中需要处理多组列对的需求，循环调用style.apply()即可：

# 定义需要处理的所有Current列
target_current_cols = [("A", "Current"), ("B", "Current")]
df_style = df.style

for current_col in target_current_cols:
    # 从Current列推导对应的Prev列
    prev_col = (current_col[0], "Prev")
    df_style = df_style.apply(
        movement,
        current_col=current_col,
        prev_col=prev_col,
        color_one="green",
        color_two="red",
        original_df=df,
        axis=None,
        subset=[current_col],
    )

df_style.render()

关键提示 #

• 多索引列的子集还可以用pd.IndexSlice更清晰地指定，比如subset=pd.IndexSlice[:, ("A", "Current")]，效果和[("A", "Current")]一致。
• 用subset时一定要注意：传入函数的是子集化后的DataFrame，函数内不能访问子集外的列，除非额外传入原DataFrame。

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