Graph Coloring in Python

color.py

# Author: Guillem G. Subies <guillemggsubies@gmail.com>
#
# License: MIT


import numpy as np


def color_graph(edges, extra_iterations=5):
    """Welsh Powell (Greedy) Algorith to color graphs + recoloring at the end.
    
    Parameters
    ----------
    edges : list of tuples of shape (n_edges, )
        List of of edges.
    extra_iterations : int
        How many iterations of recoloring will the algorithm do.

    Returns
    -------
    solution : list of shape (n_points, )
        List of colors (numbers) for every point in the graph.
    """

    edges = np.asarray(edges)
    vertex_count = np.amax(edges) + 1

    # Create the graph matrix
    graph = np.zeros(shape=(vertex_count, vertex_count), dtype=np.int8)
    rows, cols = zip(*edges)
    graph[rows, cols] = 1
    graph[cols, rows] = 1
    del rows, cols, edges

    # Gets the degree of every vertex and sorts them in a descending order
    colsum = np.argsort(-graph.sum(axis=0), kind="stable")

    colors = np.full(len(graph), -1, dtype=np.int8)
    for index in colsum:
        for color in range(vertex_count):
            for i, node in enumerate(graph[index]):
                if node == 1 and colors[i] == color and i != index:
                    break
            else:
                colors[index] = color
                break

    # Recolor TODO: Use Kemp Chains
    for p in range(extra_iterations):
        max_colores = colors.max()
        for index in colsum:
            adjacent_colors = {
                colors[i] for i, node in enumerate(graph[index]) if node == 1
            }
            for color in range(max_colores):
                if (
                    color != colors[index]
                    and color not in adjacent_colors
                    and color != max_colores
                ):
                    colors[index] = color
                    break

    return colors