Add Misra-Gries edge coloring method

docs/source/api.rst

@@ -198,6 +198,7 @@ Other Algorithm Functions
    rustworkx.transitivity
    rustworkx.core_number
    rustworkx.graph_greedy_color
+   rustworkx.graph_misra_gries_edge_color
    rustworkx.metric_closure
    rustworkx.is_planar
 

src/coloring.rs

@@ -10,12 +10,20 @@
 // License for the specific language governing permissions and limitations
 // under the License.
 
-use crate::graph;
+use crate::{graph, StablePyGraph};
 use rustworkx_core::coloring::greedy_node_color;
+use std::collections::HashSet;
 
+use petgraph::graph::{EdgeIndex, NodeIndex};
+use petgraph::visit::EdgeRef;
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
 use pyo3::Python;
+use rustworkx_core::dictmap::DictMap;
+use rustworkx_core::dictmap::*;
+
+use petgraph::visit::IntoEdgeReferences;
+use petgraph::Undirected;
 
 /// Color a :class:`~.PyGraph` object using a greedy graph coloring algorithm.
 ///
@@ -59,3 +67,209 @@ pub fn graph_greedy_color(py: Python, graph: &graph::PyGraph) -> PyResult<PyObje
     }
     Ok(out_dict.into())
 }
+
+struct MisraGriesAlgorithm<'a> {
+    graph: &'a StablePyGraph<Undirected>,
+    colors: DictMap<EdgeIndex, usize>,
+}
+
+impl<'a> MisraGriesAlgorithm<'a> {
+    pub fn new(graph: &'a StablePyGraph<Undirected>) -> Self {
+        let colors: DictMap<EdgeIndex, usize> = DictMap::new();
+        MisraGriesAlgorithm { graph, colors }
+    }
+
+    // Computes colors used at node u
+    fn get_used_colors(&self, u: NodeIndex) -> Vec<usize> {
+        let mut used_colors: Vec<usize> = Vec::new();
+        for edge in self.graph.edges(u) {
+            if let Some(color) = self.colors.get(&edge.id()) {
+                used_colors.push(*color);
+            }
+        }
+        used_colors
+    }
+
+    // Returns the smallest free (aka unused) color at node u
+    fn get_free_color(&self, u: NodeIndex) -> usize {
+        let used_colors = self.get_used_colors(u);
+        let mut c: usize = 0;
+        while used_colors.contains(&c) {
+            c += 1;
+        }
+        c
+    }
+
+    // Returns if color c is free at node u
+    fn is_free_color(&self, u: NodeIndex, c: usize) -> bool {
+        let used_colors = self.get_used_colors(u);
+        !used_colors.contains(&c)
+    }
+
+    // Returns the maximal fan on edge ee = (u, v) at u
+    fn get_maximal_fan(
+        &self,
+        ee: EdgeIndex,
+        u: NodeIndex,
+        v: NodeIndex,
+    ) -> Vec<(EdgeIndex, NodeIndex)> {
+        let mut fan: Vec<(EdgeIndex, NodeIndex)> = Vec::new();
+        fan.push((ee, v));
+
+        let mut neighbors: Vec<(EdgeIndex, NodeIndex)> = Vec::new();
+        for edge in self.graph.edges(u) {
+            neighbors.push((edge.id(), edge.target()));
+        }
+
+        let mut last_node = v;
+        let position_v = neighbors.iter().position(|x| x.1 == v).unwrap();
+        neighbors.remove(position_v);
+
+        let mut fan_extended: bool = true;
+        while fan_extended {
+            fan_extended = false;
+
+            for (edge_index, z) in &neighbors {
+                if let Some(color) = self.colors.get(edge_index) {
+                    if self.is_free_color(last_node, *color) {
+                        fan_extended = true;
+                        last_node = *z;
+                        fan.push((*edge_index, *z));
+                        let position_z = neighbors.iter().position(|x| x.1 == *z).unwrap();
+                        neighbors.remove(position_z);
+                        break;
+                    }
+                }
+            }
+        }
+
+        fan
+    }
+
+    fn flip_color(&self, c: usize, d: usize, e: usize) -> usize {
+        if e == c {
+            d
+        } else {
+            c
+        }
+    }
+
+    // Returns the longest path starting at node u with alternating colors c, d, c, d, c, etc.
+    fn get_cdu_path(&self, u: NodeIndex, c: usize, d: usize) -> Vec<(EdgeIndex, usize)> {
+        let mut path: Vec<(EdgeIndex, usize)> = Vec::new();
+        let mut cur_node: NodeIndex = u;
+        let mut cur_color = c;
+        let mut path_extended = true;
+
+        while path_extended {
+            path_extended = false;
+            for edge in self.graph.edges(cur_node) {
+                if let Some(color) = self.colors.get(&edge.id()) {
+                    if *color == cur_color {
+                        path_extended = true;
+                        path.push((edge.id(), cur_color));
+                        cur_node = edge.target();
+                        cur_color = self.flip_color(c, d, cur_color);
+                        break;
+                    }
+                }
+            }
+        }
+        path
+    }
+
+    fn check_coloring(&self) -> bool {
+        for edge in self.graph.edge_references() {
+            match self.colors.get(&edge.id()) {
+                Some(_color) => (),
+                None => {
+                    println!("Problem edge {:?} has no color assigned", edge);
+                    return false;
+                }
+            }
+        }
+
+        let mut max_color = 0;
+        for node in self.graph.node_indices() {
+            let mut used_colors: HashSet<usize> = HashSet::new();
+            let mut num_edges = 0;
+            for edge in self.graph.edges(node) {
+                num_edges += 1;
+                match self.colors.get(&edge.id()) {
+                    Some(color) => {
+                        used_colors.insert(*color);
+                        if max_color < *color {
+                            max_color = *color;
+                        }
+                    }
+                    None => {
+                        println!("Problem: edge {:?} has no color assigned", edge);
+                        return false;
+                    }
+                }
+            }
+            if used_colors.len() < num_edges {
+                println!("Problem: node {:?} does not have enough colors", node);
+                return false;
+            }
+        }
+
+        println!("Coloring is OK, max_color = {}", max_color);
+        true
+    }
+
+    pub fn run_algorithm(&mut self) -> &DictMap<EdgeIndex, usize> {
+        for edge in self.graph.edge_references() {
+            let u: NodeIndex = edge.source();
+            let v: NodeIndex = edge.target();
+            let fan = self.get_maximal_fan(edge.id(), u, v);
+            let c = self.get_free_color(u);
+            let d = self.get_free_color(fan.last().unwrap().1);
+
+            // find cdu-path
+            let cdu_path = self.get_cdu_path(u, d, c);
+
+            // invert colors on cdu-path
+            for (edge_index, color) in cdu_path {
+                let flipped_color = self.flip_color(c, d, color);
+                self.colors.insert(edge_index, flipped_color);
+            }
+
+            // find sub-fan fan[0..w] such that d is free on fan[w]
+            let mut w = 0;
+            for (i, (_, z)) in fan.iter().enumerate() {
+                if self.is_free_color(*z, d) {
+                    w = i;
+                    break;
+                }
+            }
+
+            // rotate fan
+            for i in 1..w + 1 {
+                let next_color = self.colors.get(&(fan[i].0)).unwrap();
+                self.colors.insert(fan[i - 1].0, *next_color);
+            }
+
+            // fill additional color
+            self.colors.insert(fan[w].0, d);
+        }
+
+        self.check_coloring();
+
+        &self.colors
+    }
+}
+
+#[pyfunction]
+#[pyo3(text_signature = "(graph, /)")]
+pub fn graph_misra_gries_edge_color(py: Python, graph: &graph::PyGraph) -> PyResult<PyObject> {
+    let mut mg = MisraGriesAlgorithm::new(&graph.graph);
+
+    let colors = mg.run_algorithm();
+
+    let out_dict = PyDict::new(py);
+    for (edge, color) in colors {
+        out_dict.set_item(edge.index(), color)?;
+    }
+    Ok(out_dict.into())
+}

src/lib.rs

@@ -433,6 +433,7 @@ fn rustworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(graph_astar_shortest_path))?;
     m.add_wrapped(wrap_pyfunction!(digraph_astar_shortest_path))?;
     m.add_wrapped(wrap_pyfunction!(graph_greedy_color))?;
+    m.add_wrapped(wrap_pyfunction!(graph_misra_gries_edge_color))?;
     m.add_wrapped(wrap_pyfunction!(graph_tensor_product))?;
     m.add_wrapped(wrap_pyfunction!(digraph_tensor_product))?;
     m.add_wrapped(wrap_pyfunction!(directed_gnp_random_graph))?;

tests/rustworkx_tests/graph/test_coloring.py

@@ -44,3 +44,47 @@ def test_simple_graph_large_degree(self):
         graph.add_edge(node_a, node_c, 1)
         res = rustworkx.graph_greedy_color(graph)
         self.assertEqual({0: 0, 1: 1, 2: 1}, res)
+
+
+class TestMisraGriesColoring(unittest.TestCase):
+
+    def test_simple_graph(self):
+        graph = rustworkx.PyGraph()
+        node0 = graph.add_node(0)
+        node1 = graph.add_node(1)
+        node2 = graph.add_node(2)
+        node3 = graph.add_node(3)
+
+        graph.add_edge(node0, node1, 1)
+        graph.add_edge(node0, node2, 1)
+        graph.add_edge(node1, node2, 1)
+        graph.add_edge(node2, node3, 1)
+
+        print("======================")
+        res = rustworkx.graph_misra_gries_edge_color(graph)
+        print(f"{res = }")
+
+
+    def test_grid(self):
+        graph = rustworkx.generators.grid_graph(10, 10);
+
+        print("======================")
+        res = rustworkx.graph_misra_gries_edge_color(graph)
+        print(f"{res = }")
+
+
+    def test_heavy_hex(self):
+        graph = rustworkx.generators.heavy_hex_graph(7);
+
+        print("======================")
+        res = rustworkx.graph_misra_gries_edge_color(graph)
+        print(f"{res = }")
+
+    def test_barbell(self):
+
+        # graph = rustworkx.generators.barbell_graph(4, 3)
+        graph = rustworkx.generators.complete_graph(10)
+
+        print("======================")
+        res = rustworkx.graph_misra_gries_edge_color(graph)
+        print(f"{res = }")