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