Skip to main content

weavatrix_graph/algo/
coloring.rs

1use crate::IndexUndirectedGraphView;
2use crate::Vec;
3use alloc::collections::{BTreeSet, BinaryHeap};
4use core::cmp::Reverse;
5
6type Indexed<Node> = (Vec<Option<Node>>, Vec<Vec<usize>>);
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct Coloring<Node> {
10    assignments: Vec<(Node, usize)>,
11    color_count: usize,
12}
13
14impl<Node> Coloring<Node> {
15    #[must_use]
16    pub fn assignments(&self) -> &[(Node, usize)] {
17        &self.assignments
18    }
19
20    #[must_use]
21    pub const fn color_count(&self) -> usize {
22        self.color_count
23    }
24}
25
26/// Produces a deterministic DSATUR coloring.
27///
28/// DSATUR is exact for several common graph classes but is a heuristic for the
29/// general NP-hard minimum-coloring problem.
30pub fn dsatur_coloring<G>(graph: &G) -> Option<Coloring<G::Node>>
31where
32    G: IndexUndirectedGraphView,
33{
34    let (nodes, adjacency) = indexed(graph)?;
35    let mut colors = vec![None; graph.node_bound()];
36    let mut adjacent_colors = vec![BTreeSet::new(); graph.node_bound()];
37    let mut queue = BinaryHeap::new();
38    for (slot, node) in nodes.iter().enumerate() {
39        if node.is_some() {
40            queue.push((0, adjacency[slot].len(), Reverse(slot)));
41        }
42    }
43    while let Some((saturation, _, Reverse(node))) = queue.pop() {
44        if colors[node].is_some() || saturation != adjacent_colors[node].len() {
45            continue;
46        }
47        let color = (0..=adjacency.len())
48            .find(|color| !adjacent_colors[node].contains(color))
49            .unwrap_or(0);
50        colors[node] = Some(color);
51        for &neighbor in &adjacency[node] {
52            if colors[neighbor].is_none() && adjacent_colors[neighbor].insert(color) {
53                queue.push((
54                    adjacent_colors[neighbor].len(),
55                    adjacency[neighbor].len(),
56                    Reverse(neighbor),
57                ));
58            }
59        }
60    }
61    let assignments = nodes
62        .into_iter()
63        .enumerate()
64        .filter_map(|(slot, node)| Some((node?, colors[slot]?)))
65        .collect::<Vec<_>>();
66    let color_count = assignments
67        .iter()
68        .map(|(_, color)| color + 1)
69        .max()
70        .unwrap_or(0);
71    Some(Coloring {
72        assignments,
73        color_count,
74    })
75}
76
77fn indexed<G: IndexUndirectedGraphView>(graph: &G) -> Option<Indexed<G::Node>> {
78    let mut nodes = vec![None; graph.node_bound()];
79    let mut adjacency = vec![Vec::new(); graph.node_bound()];
80    for node in graph.node_indices() {
81        let slot = G::node_slot(node);
82        nodes[slot] = Some(node);
83        adjacency[slot] = graph
84            .incident_edges(node)
85            .filter_map(|edge| graph.opposite(edge, node))
86            .map(G::node_slot)
87            .collect();
88        if adjacency[slot].contains(&slot) {
89            return None;
90        }
91        adjacency[slot].sort_unstable();
92        adjacency[slot].dedup();
93    }
94    Some((nodes, adjacency))
95}