如何在Matplotlib中识别并分组相似颜色的点？

你当前的代码是按每行数据连接点，而目标是将颜色相似的点跨行连接并分组展示。要实现这个效果，核心思路是先对颜色进行聚类，把相似颜色归为同一组，再按组连接对应的数据点。

示例数据 #

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

x = np.array([0,1,2,3,4])
y = np.array([[0.0, 1.0, 1.4, 3.0, 4.0],
              [1.0, 1.4, 1.9, 2.1, 3.0],
              [0.9, 1.0, 1.4, 3.0, 3.1],
              [0.3, 1.0, 1.4, 3.0, 3.7],
              [0.1, 1.0, 1.4, 3.0, 3.9]])

color = [[1,0,0.02], [0.01,.91,0], [0,1,0], [0,0.04,1], [1,.91,0],
         [0,1,0.01], [0.00,1,0], [1,.91,0], [1,0.01,0], [0,0,1],
         [1,.91,0.00], [0.03,1,0], [0,1,0], [0,0.00,1], [1,0,0],
         [1,.91,0.03], [0.05,1,0], [0,1,0], [0,0.03,1], [1,0,0],
         [1,.91,0.00], [0.01,.91,0], [0,1,0], [0,0.03,1], [1,0,0]]

实现步骤 #

1. 预处理数据与颜色：将颜色列表转为numpy数组，同时把x、y数据展开为一维数组，方便对应每个点的坐标和颜色。
2. 颜色聚类：用KMeans将相似颜色聚为一类，聚类数量可根据颜色种类调整（示例中颜色大致分为红、绿、蓝、黄4类）。
3. 按组绘制连线：对每个聚类组，收集该组所有点的(x,y)坐标，按x值排序后绘制连线，保证连线顺序正确。
4. 绘制散点图：保留原散点的颜色，确保视觉一致性。

完整代码 #

# 数据预处理
color_np = np.array(color)
x_flat = np.repeat(x, y.shape[1])
y_flat = y.flatten()

# 颜色聚类（根据实际颜色种类调整n_clusters）
kmeans = KMeans(n_clusters=4, random_state=42)
cluster_labels = kmeans.fit_predict(color_np)

# 创建绘图
plt.figure(figsize=(8,6))

# 按聚类组绘制连线
for label in np.unique(cluster_labels):
    # 获取当前组的所有点
    mask = cluster_labels == label
    group_x = x_flat[mask]
    group_y = y_flat[mask]
    
    # 按x值排序，保证连线顺序正确
    sort_idx = np.argsort(group_x)
    sorted_x = group_x[sort_idx]
    sorted_y = group_y[sort_idx]
    
    # 绘制组内连线，颜色用聚类中心的颜色（或组内任意颜色）
    plt.plot(sorted_x, sorted_y, color=kmeans.cluster_centers_[label], lw=0.7)

# 绘制散点图，保留原颜色
plt.scatter(x_flat, y_flat, s=12, c=color)

plt.show()

说明 #

• 聚类数量n_clusters可根据实际颜色的相似度调整，如果颜色区分度高，可适当增加数量；如果有相近颜色，可减少数量。
• 连线颜色使用聚类中心的颜色，保证组内连线颜色与点的颜色风格一致。
• 对每组点按x值排序，避免连线出现交叉混乱的情况。

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