如何在matplotlib/seaborn热图中为每个单元格条件定义颜色

你的两个需求均可以实现，以下是具体代码实现：

1. 按条件为特定数值单元格染色 #

你可以通过坐标定位的方式，给满足条件的单元格单独设置颜色，比如将数值为60的单元格设为绿色：

import seaborn as sns
import numpy as np
import matplotlib.pyplot as plt

# 替换为你的实际数据
data = np.random.randint(0, 100, size=(15, 15)) 

# 绘制基础热图
ax = sns.heatmap(data, cmap="coolwarm", annot=True)

# 筛选所有数值等于60的单元格
mask = data == 60

# 批量为符合条件的单元格上色
for i in range(data.shape[0]):
    for j in range(data.shape[1]):
        if mask[i, j]:
            ax.add_patch(plt.Rectangle((j, i), 1, 1, color="green", edgecolor="white"))

plt.show()

2. 传入与数据同形状的numpy颜色数组绘图 #

完全支持该操作，你可以提前生成和数据形状严格一致的RGBA颜色数组，绘图时直接传入即可：

from matplotlib.colors import ListedColormap

# 生成形状为(行数, 列数, 4)的颜色数组，最后一维对应RGBA四个通道（取值范围0~1）
color_arr = np.zeros(data.shape + (4,))

# 自定义颜色规则，可根据需求任意调整
for i in range(data.shape[0]):
    for j in range(data.shape[1]):
        val = data[i, j]
        if val == 60:
            color_arr[i,j] = [0, 1, 0, 1] # 绿色，不透明
        elif val < 30:
            color_arr[i,j] = [0.8, 0.9, 1, 1] # 浅蓝
        elif val < 70:
            color_arr[i,j] = [0.2, 0.5, 0.9, 1] # 深蓝
        else:
            color_arr[i,j] = [0.9, 0.2, 0.2, 1] # 红色

# 传入自定义颜色数组绘图
ax = sns.heatmap(
    data, 
    annot=True, # 保留原始数值标注
    cmap=ListedColormap(color_arr.reshape(-1, 4)),
    cbar=False # 如果不需要色条可以关闭，需要的话建议自定义离散色阶匹配规则
)
plt.show()

效率优化提示 #

如果数据尺寸较大，不需要循环遍历生成颜色数组，可以批量操作：

from matplotlib.cm import get_cmap

# 用基础色阶生成默认颜色
norm = plt.Normalize(data.min(), data.max())
base_cmap = get_cmap("YlGnBu")
color_arr = base_cmap(norm(data))

# 批量替换符合条件的单元格颜色
color_arr[data == 60] = [0, 1, 0, 1]

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