Python实现n个节点分组自动分配唯一绘图颜色

问题背景 #

现有存储节点分组标签的数组X，数组索引对应节点编号，元素值为该节点所属的分组ID。绘图时需要为同一分组的所有节点分配唯一可区分的颜色，原有硬编码if-else的写法仅支持固定3个分组，无法适配分组数动态变化的场景。

测试用的分组数组如下：

X = [1, 0, 2, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]

原有硬编码实现：

color_map = []
for i in X:
    if i==0:
        color_map.append('blue')
    elif i==1:
        color_map.append('yellow')
    else: 
        color_map.append('green')

通用实现代码 #

核心逻辑是先自动提取所有唯一分组ID，建立分组到颜色的映射字典，再批量生成每个节点对应的颜色，不需要手动判断分组数量。

方案1：基于matplotlib色卡（推荐，区分度最高） #

直接调用matplotlib内置的校准色板，颜色区分度有保证，支持几个到上百个分组的场景，不需要自己维护颜色列表：

import matplotlib.pyplot as plt

X = [1, 0, 2, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]

# 自动提取所有唯一分组，统计分组总数
unique_groups = sorted(set(X))
group_count = len(unique_groups)

# 选择色板：分组数≤20用tab20（高区分度离散色），超过20可替换为hsv/gist_ncar
cmap = plt.get_cmap('tab20', group_count)
# 生成分组-颜色映射字典
group_to_color = {gid: cmap(idx) for idx, gid in enumerate(unique_groups)}
# 生成每个节点对应的颜色列表
color_map = [group_to_color[gid] for gid in X]

• 分组数不超过20时，tab20色板的所有颜色都是人工校准的高区分度色彩，不会出现相邻分组颜色相近难以分辨的问题
• 分组数超过20时，把色板参数替换为'hsv'即可支持最多上百个分组的颜色区分

方案2：兼容原有命名颜色的实现 #

如果需要沿用之前的命名颜色（blue/yellow/green这类字符串格式），可以预置足够长的命名颜色列表，分组数超出预置范围时自动补随机RGB颜色：

import random

# 预置高区分度命名颜色，可根据需要自行扩充
PRESET_COLORS = [
    'blue', 'yellow', 'green', 'red', 'purple', 'orange', 'pink', 'brown',
    'gray', 'olive', 'cyan', 'magenta', 'teal', 'navy', 'maroon', 'lime'
]

X = [1, 0, 2, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
unique_groups = sorted(set(X))

group_to_color = {}
for idx, gid in enumerate(unique_groups):
    if idx < len(PRESET_COLORS):
        group_to_color[gid] = PRESET_COLORS[idx]
    else:
        # 预置颜色不足时生成随机RGB颜色
        group_to_color[gid] = (random.random(), random.random(), random.random())

color_map = [group_to_color[gid] for gid in X]

这个方案和原有硬编码逻辑的输出完全兼容，分组数为3时，输出的颜色和之前写的blue/yellow/green完全一致。

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