petgraph/algo/
coloring.rs

1use alloc::{collections::BinaryHeap, vec};
2use core::hash::Hash;
3
4use hashbrown::{HashMap, HashSet};
5
6use crate::scored::MaxScored;
7use crate::visit::{IntoEdges, IntoNodeIdentifiers, NodeIndexable, VisitMap, Visitable};
8
9/// \[Generic\] [DStatur algorithm][1] to properly color a non weighted undirected graph.
10///
11///
12/// This is a heuristic. So, it does not necessarily return a minimum coloring.
13/// The graph must be undirected. It should not contain loops.
14///
15/// # Arguments
16/// * `graph`: undirected graph without loops.
17///
18/// # Returns
19/// Returns a tuple of:
20/// * [`struct@hashbrown::HashMap`] that associates to each `NodeId` its color.
21/// * `usize`: the number of used colors.
22///
23/// # Complexity
24/// * Time complexity: **O((|V| + |E|)log(|V|)**.
25/// * Auxiliary space: **O(|V| + |E|)**.
26///
27/// where **|V|** is the number of nodes and **|E|** is the number of edges.
28///
29/// [1]: https://en.wikipedia.org/wiki/DSatur
30///
31/// # Example
32/// ```rust
33/// use petgraph::{Graph, Undirected};
34/// use petgraph::algo::dsatur_coloring;
35///
36/// let mut graph: Graph<(), (), Undirected> = Graph::new_undirected();
37/// let a = graph.add_node(());
38/// let b = graph.add_node(());
39/// let c = graph.add_node(());
40/// let d = graph.add_node(());
41/// let e = graph.add_node(());
42/// let f = graph.add_node(());
43///
44/// graph.extend_with_edges(&[
45///     (a, b),
46///     (b, c),
47///     (c, d),
48///     (d, e),
49///     (e, f),
50///     (f, a),
51/// ]);
52///
53/// // a ----- b ----- c
54/// // |               |
55/// // |               |
56/// // |               |
57/// // f ----- e------ d
58///
59/// let (coloring, nb_colors) = dsatur_coloring(&graph);
60/// assert_eq!(nb_colors, 2);
61/// assert_ne!(coloring[&a], coloring[&b]);
62/// ```
63pub fn dsatur_coloring<G>(graph: G) -> (HashMap<G::NodeId, usize>, usize)
64where
65    G: IntoEdges + IntoNodeIdentifiers + Visitable + NodeIndexable,
66    G::NodeId: Eq + Hash,
67{
68    let ix = |v| graph.to_index(v);
69    let n = graph.node_bound();
70
71    let mut degree_map = vec![0; n];
72    let mut queue = BinaryHeap::with_capacity(n);
73    let mut colored = HashMap::with_capacity(n);
74    let mut adj_color_map = vec![HashSet::new(); n];
75    let mut seen = graph.visit_map();
76    let mut max_color = 0;
77
78    for node in graph.node_identifiers() {
79        let degree = graph.edges(node).count();
80        queue.push(MaxScored((0, degree), node));
81        degree_map[ix(node)] = degree;
82    }
83
84    while let Some(MaxScored(_, node)) = queue.pop() {
85        if seen.is_visited(&node) {
86            continue;
87        }
88        seen.visit(node);
89
90        let adj_color = &adj_color_map[ix(node)];
91        let mut color = 0;
92        while adj_color.contains(&color) {
93            color += 1;
94        }
95
96        colored.insert(node, color);
97        max_color = max_color.max(color);
98
99        for nbor in graph.neighbors(node) {
100            if let Some(adj_color) = adj_color_map.get_mut(ix(nbor)) {
101                adj_color.insert(color);
102                queue.push(MaxScored((adj_color.len(), degree_map[ix(nbor)]), nbor));
103            }
104        }
105    }
106
107    (colored, max_color + 1)
108}