Pandas样式突出显示对角线和非对角线元素

import pandas as pd
import numpy as np

# 创建示例DataFrame
df = pd.DataFrame(np.random.randint(0, 10, size=(5, 5)), columns=['A', 'B', 'C', 'D', 'E'])

def highlight_diagonal(x):
    """
    根据元素位置返回不同的背景颜色
    """
    diagonal_color = 'yellow'  # 对角线元素的颜色
    non_diagonal_color = 'lightblue'  # 非对角线元素的颜色

    # 创建一个空的DataFrame样式对象
    style = pd.DataFrame('', index=x.index, columns=x.columns)

    # 对角线元素
    np.fill_diagonal(style.values, f'background-color: {diagonal_color}')

    # 非对角线元素
    style.mask(np.eye(len(x), dtype=bool), f'background-color: {non_diagonal_color}', inplace=True)

    return style

# 应用样式函数
highlighted_df = df.style.apply(highlight_diagonal, axis=None)

# 展示结果
highlighted_df