如何在Matplotlib中绘制按currState变色的时间-空间水平条形图

我明白你想要实现的效果——在单条水平条形图上，根据不同时间区间内的currState值，对应显示红、黄、绿三种颜色。你的现有代码还没处理颜色分段的核心逻辑，接下来我一步步帮你搞定这个需求。

核心思路 #

要实现这个效果，我们不能直接绘制一整根条形，而是需要：

1. 把连续的时间序列转换成时间差数值（方便计算条形长度）
2. 把连续相同状态的时间段合并，得到每个颜色段的起始位置和持续时长
3. 循环绘制每个分段的条形，对应各自的颜色

完整代码实现 #

首先假设你已经导入了必要的库：import pandas as pd、import matplotlib.pyplot as plt、import numpy as np

# 你的颜色映射字典
category_colors = {'red': [2,3], 'yellow': [5,6], 'green': [7,8]}

# 1. 数据预处理：先把timestamp转为datetime类型，计算时间差
df_sample['timestamp'] = pd.to_datetime(df_sample['timestamp'])
# 计算每个时间点相对于起始时间的秒数（用于条形图的位置和宽度）
df_sample['time_diff'] = (df_sample['timestamp'] - df_sample['timestamp'].iloc[0]).dt.total_seconds()

# 2. 标记连续相同状态的分组，聚合得到每个分段的信息
df_sample['group'] = (df_sample['currState'] != df_sample['currState'].shift()).cumsum()
grouped = df_sample.groupby('group').agg(
    start_time=('time_diff', 'first'),
    duration=('time_diff', lambda x: x.iloc[-1] - x.iloc[0]),
    state=('currState', 'first')
).reset_index()

# 3. 转换颜色映射：从颜色对应状态列表，转为状态值对应颜色
state_to_color = {}
for color, states in category_colors.items():
    for s in states:
        state_to_color[s] = color

# 4. 绘制分段条形图
fig, ax = plt.subplots(figsize=(10, 1))
y_pos = 0  # 只有一个路口，y轴位置固定
intersection_id = df_sample['IntersectionId'].iloc[0]

# 循环每个分段，绘制对应颜色的条形
for _, row in grouped.iterrows():
    bar_color = state_to_color[row['state']]
    # barh参数：y轴位置，条形宽度（持续时长），条形起始位置（起始时间差）
    ax.barh(y_pos, row['duration'], left=row['start_time'], color=bar_color)

# 设置坐标轴和标题
ax.set_yticks([y_pos])
ax.set_yticklabels([intersection_id])
ax.set_xlabel('timestamp')
ax.set_ylabel('space')
ax.set_title('TimeSpace')

# 优化x轴显示：把时间差转回原始时间戳，旋转避免重叠
start_time = df_sample['timestamp'].iloc[0]
xticks = ax.get_xticks()
ax.set_xticklabels([start_time + pd.Timedelta(seconds=t) for t in xticks], rotation=45)

# 添加自定义图例
from matplotlib.patches import Patch
legend_elements = [
    Patch(facecolor=color, label=f'States {states}') 
    for color, states in category_colors.items()
]
ax.legend(handles=legend_elements, loc='upper right')

plt.tight_layout()
plt.show()

代码解释 #

• 数据预处理：把字符串类型的timestamp转为datetime，计算时间差是为了让Matplotlib能处理数值型的条形宽度和位置
• 分组连续状态：通过shift()对比当前行和上一行的状态，标记分组后聚合，得到每个颜色段的起始时间、持续时长和对应状态
• 颜色映射转换：把原来的颜色→状态列表字典反转成状态值→颜色，方便快速查找每个状态对应的颜色
• 分段绘制：循环每个分组，用barh()绘制分段条形，left参数控制条形的起始位置，width控制长度，这样就能在同一y轴位置上拼接不同颜色的分段

效果说明 #

运行这段代码后，你会得到一条水平条形图，不同时间区间会根据currState的值自动切换为红、黄、绿三色，完全符合你的期望效果。

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