如何基于4种条件使用Matplotlib绘制散点图？

核心思路 #

放弃嵌套循环遍历行的低效方式，利用pandas布尔索引直接筛选符合条件的数据子集，结合matplotlib批量绘制散点，既高效又易维护。

代码实现 #

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

# 定义4种条件对应的样式：(筛选条件, 颜色, 标记符号)
plot_configs = [
    ((data['t'] == 1) & (data['p'] == 1), 'green', 'o'),
    ((data['t'] == 1) & (data['p'] == 0), 'green', 'x'),
    ((data['t'] == 0) & (data['p'] == 1), 'red', 'x'),
    ((data['t'] == 0) & (data['p'] == 0), 'red', 'o'),
]

# 创建画布并循环绘制
plt.figure(figsize=(10, 6))
for condition, color, marker in plot_configs:
    # 筛选符合当前条件的数据
    subset = data[condition]
    # 绘制散点，alpha可选，缓解点重叠问题
    plt.scatter(subset['validation_WES'], subset['CalcESR_R3_refitted'],
                color=color, marker=marker, alpha=0.6)

# 设置图表标签与自定义图例
plt.xlabel('validation_WES')
plt.ylabel('CalcESR_R3_refitted')
plt.title('Multi-condition Scatter Plot')

# 手动创建图例项（因为散点是分批次绘制的）
legend_elements = [
    Line2D([0], [0], color='green', marker='o', linestyle='', markersize=10, label='t=1, p=1'),
    Line2D([0], [0], color='green', marker='x', linestyle='', markersize=10, label='t=1, p=0'),
    Line2D([0], [0], color='red', marker='x', linestyle='', markersize=10, label='t=0, p=1'),
    Line2D([0], [0], color='red', marker='o', linestyle='', markersize=10, label='t=0, p=0'),
]
plt.legend(handles=legend_elements)

plt.show()

为什么比嵌套循环好 #

1. 效率高：pandas的布尔索引是基于向量化操作实现的，远快于逐行遍历的循环；
2. 代码简洁：用统一的配置列表管理所有条件和样式，后期修改或新增条件只需调整列表即可；
3. 易调试：每一步的筛选逻辑清晰，便于排查问题。

内容的提问来源于stack exchange，提问作者all.m