A collection of graph coloring algorithms

graphcoloring.py

"""A collection of graph coloring algorithms

The problem of assigning colors to the vertices of a graph such that no two adjacent vertices have the same color.
"""
import random
import typing
from collections import defaultdict

import networkx


T = typing.TypeVar("T")


def greedy_algorithm(G: networkx.Graph, colorset: set[T]) -> dict[int, T] | None:
    """A coloring algorithm that assigns colors starting from a generic random vertex"""
    if not G or not colorset: return None

    coloring = {}

    for node in G:        
        unavailable = set()
        
        for neighbor in G[node]:
            if neighbor in coloring:    
                unavailable.add(coloring[neighbor])
        
        labelset = colorset - unavailable
        if not labelset: return None
        
        coloring[node] = labelset.pop()
 
    return coloring


def backtracking(G: networkx.Graph, colorset: set[T]) -> dict[int, T] | None:
    """A coloring algorithm that assigns colors starting from a specific vertex
    
    It recursively backtracks checking for assignment safety
    """
    if not G or not colorset: return None
    
    coloring = {}
    
    def available(node: int, color: T) -> bool:
        unavailable = set()

        for neighbor in G[node]:
            if neighbor in coloring:               
                unavailable.add(coloring[neighbor])

        return color not in unavailable
    
    def f(node: int) -> bool:
        nonlocal coloring
        
        if node == len(G): return True
        
        for color in colorset:
            if available(node, color):
                coloring[node] = color
                
                if f(node + 1): return True

                del coloring[node]

    return coloring if f(0) else None


def dsatur(G: networkx.Graph, colorset: set[T]) -> dict[int, T] | None:
    """The DSatur coloring algorithm [1]_
    
    A heuristic algorithm that assigns colors starting from the vertex with the highest degree of saturation.
    
    Notes
    -----
    The saturation is the number of different colors used by adjacent vertices.

    References
    ----------
    .. [1] https://doi.org/10.1145/359094.359101
       New Methods to Color the Vertices of a Graph
    """
    if not G or not colorset: return None

    coloring = {}
    colorset = {color: 0 for color in colorset}

    def degree(node: int) -> int:
        return len(G[node])

    def saturation(node: int) -> int:
        return len(set(coloring[neighbor] for neighbor in G[node] if neighbor in coloring))

    def available(node: int, color: T) -> bool:
        unavailable = set()

        for neighbor in G[node]:
            if neighbor in coloring:               
                unavailable.add(coloring[neighbor])

        return color not in unavailable

    def assignment(node: int, color: T) -> None:
        nonlocal coloring, colorset

        coloring[ node] = color
        colorset[color] = colorset[color] + 1

    def selection() -> int | None:
        unpainted = [node for node in G if node not in coloring]
        
        if not unpainted: return None
        
        reference = defaultdict(lambda: [])
        
        for node in unpainted:
            reference[saturation(node)].append(node)
        
        max_saturation = max(reference.keys())
        
        candidateset = reference[max_saturation]

        if len(candidateset) == 1: return candidateset[0]
        
        return max(candidateset, key=degree)

    assignment(random.choice(range(len(G))), random.choice(colorset.keys()))

    node = selection()

    while node is not None:
        unpainted = [color for color in colorset if available(node, color)]

        if not unpainted: return None

        assignment(node, random.choice(unpainted))
        
        node = selection()
    
    return coloring