如何基于pandas DataFrame多日期列和指定条件生成带标注热图

以下代码可完整解决你的需求，覆盖条件判断、分类值生成、热图标注三个环节：

步骤1：导入依赖&预处理数据 #

首先统一日期格式，提取排序后的全量日期作为x轴：

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

# 原始数据
data = {'ID':['001', '002', '003', '004', '005'],
    'Aim2 Date':['9/3/19', '9/3/19', '9/13/19', '9/20/19', '9/23/19'],
    'Aim2 C/I' :['Intervention', 'Intervention', 'Intervention', 'Control', 'Intervention'],
    'Aim3 Date':['9/17/19','9/19/19','9/27/19','9/30/19','10/8/19'],
    'Aim3 C/I' :['Control', 'Intervention', 'Control', 'Withdraw', 'Withdraw']}
df = pd.DataFrame(data)

# 转换日期列为datetime格式，避免排序错误
df['Aim2 Date'] = pd.to_datetime(df['Aim2 Date'])
df['Aim3 Date'] = pd.to_datetime(df['Aim3 Date'])

# 提取所有唯一日期，按时间升序排列作为热图x轴
all_dates = pd.concat([df['Aim2 Date'], df['Aim3 Date']]).unique()
all_dates_sorted = sorted(all_dates)
# 提取ID作为y轴
ids = df['ID'].unique()

步骤2：生成热图数值矩阵（解决条件判断、分类值生成问题） #

按规则赋值，逻辑严格匹配需求：

# 初始化全0矩阵，行是ID，列是排序后日期
heatmap_data = pd.DataFrame(0, index=ids, columns=all_dates_sorted)

# 遍历每一行原始数据，按条件赋值
for _, row in df.iterrows():
    current_id = row['ID']
    # 处理Aim2的干预标记
    if row['Aim2 C/I'] == 'Intervention':
        heatmap_data.loc[current_id, row['Aim2 Date']] = 1
    # 处理Aim3的干预标记
    if row['Aim3 C/I'] == 'Intervention':
        heatmap_data.loc[current_id, row['Aim3 Date']] = 2

步骤3：绘制带标注的热图（解决热图标注问题） #

自定义配色和标注格式，可读性更高：

plt.figure(figsize=(12, 5))
# 自定义配色：0=白色，1=蓝色，2=橙色，区分度更高
cmap = ListedColormap(['white', '#1f77b4', '#ff7f0e'])

# 绘制热图，annot=True开启单元格标注，fmt='d'表示标注整数
ax = sns.heatmap(heatmap_data, 
                 annot=True, 
                 fmt='d',
                 cmap=cmap,
                 cbar_kws={'ticks': [0, 1, 2], 'label': '干预类型'},
                 xticklabels=[d.strftime('%m/%d/%y') for d in all_dates_sorted],
                 yticklabels=ids)

# 调整轴标签
ax.set_xlabel('日期')
ax.set_ylabel('用户ID')
plt.xticks(rotation=45)
plt.tight_layout()
plt.show()

补充说明 #

如果需要修改标注内容为文字而非数字，只需要将heatmap_data里的1替换为Aim2干预、2替换为Aim3干预、0替换为空字符串，同时把绘图代码里的fmt='d'改成fmt='s'即可。

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