From 8af19f0f9c3f80463e39bb691b8f13e96a5d243f Mon Sep 17 00:00:00 2001 From: Alexander Ivrii Date: Thu, 30 May 2024 14:03:04 +0300 Subject: [PATCH] Add "saturation first" and "independent set" greedy node coloring strategies (#1138) This PR adds two more standard coloring strategies to greedy node/edge coloring functions graph_greedy_color and graph_greedy_edge_color. In the literature the first strategy is known as "saturation", "DSATUR" or "SLF". We dynamically choose the vertex that has the largest number of different colors already assigned to its neighbors, and, in case of a tie, the vertex that has the largest number of uncolored neighbors. The second strategy is known as "greedy independent sets" or "GIS". We greedily find independent subsets of the graph, and assign a different color to each of these subsets. This addresses #1137, modulo the fact that there are quadrillions of different greedy node coloring algorithms, and each comes with trillion variations. Both new strategies can be combined with preset_color_fn (for node coloring). There is also a new Rust enum GreedyStrategy exposed as a python class that allows to specify which of the three currently supported greedy strategies is used. This is, calling the functions from the Python code would be as follows: import rustworkx as rx graph = rx.generators.generalized_petersen_graph(5, 2) coloring = rx.graph_greedy_color(graph, greedy_strategy=rx.GreedyStrategy.Degree) coloring = rx.graph_greedy_color(graph, greedy_strategy=rx.GreedyStrategy.Saturation) coloring = rx.graph_greedy_color(graph, greedy_strategy=rx.GreedyStrategy.IndependentSet) coloring = rx.graph_greedy_color(graph) coloring = rx.graph_greedy_edge_color(graph) coloring = rx.graph_greedy_edge_color(graph, greedy_strategy=rx.GreedyStrategy.Degree) coloring = rx.graph_greedy_edge_color(graph, greedy_strategy=rx.GreedyStrategy.Saturation) coloring = rx.graph_greedy_edge_color(graph, greedy_strategy=rx.GreedyStrategy.IndependentSet) The greedy_graph_edge_color function has been also extended to support the preset_color_fn argument (though in this case it's an optional map from an edge index to either a color or to None) * initial commit * fix and test * python formatting * creating enums * also passing greedy strategy to edge coloring functions; improving docstrings * release notes' * python test * reno fix * reno fix * minor cleanup * small rust cleanup * adding greedy node coloring using maximal independent set + reworked tests * updating top-level source, tests and docs * docs * applying fix from code review * adding IndependentSet to pyi as well * Adding a table that describes greedy strategies available * converting table to sphinx-style * fix reference to footnotes * docs improvements * another attempt * attempt to extend API with preset_color_fn for edges * minor * finally compiling * fixing pyi API and edding python tests * adding rustworkx-core tests * adding arguments to docstrings * expanding release note * another docs fix * Update releasenotes/notes/saturation-greedy-color-109d40f189590d3a.yaml Co-authored-by: Matthew Treinish * a round of renaming * restoring greedy_node_coloring API for backwards compatibility * restoring greedy_edge_coloring API for backwards compatibility * fix to rustworkx coloring * changing return type of new rustworkx-core functions to be Result * black * removing NodeIndexable in node coloring functions * and from edge coloring functions * suggestions from code review * suggestions from code review * release note update --------- Co-authored-by: Matthew Treinish --- .../api/algorithm_functions/coloring.rst | 1 + ...uration-greedy-color-109d40f189590d3a.yaml | 59 ++ rustworkx-core/src/coloring.rs | 956 +++++++++++++++++- rustworkx/__init__.pyi | 1 + rustworkx/rustworkx.pyi | 18 +- src/coloring.rs | 119 ++- src/lib.rs | 1 + tests/graph/test_coloring.py | 88 ++ 8 files changed, 1216 insertions(+), 27 deletions(-) create mode 100644 releasenotes/notes/saturation-greedy-color-109d40f189590d3a.yaml diff --git a/docs/source/api/algorithm_functions/coloring.rst b/docs/source/api/algorithm_functions/coloring.rst index c1b7c5b68..843a390cf 100644 --- a/docs/source/api/algorithm_functions/coloring.rst +++ b/docs/source/api/algorithm_functions/coloring.rst @@ -6,6 +6,7 @@ Coloring .. autosummary:: :toctree: ../../apiref + rustworkx.ColoringStrategy rustworkx.graph_greedy_color rustworkx.graph_bipartite_edge_color rustworkx.graph_greedy_edge_color diff --git a/releasenotes/notes/saturation-greedy-color-109d40f189590d3a.yaml b/releasenotes/notes/saturation-greedy-color-109d40f189590d3a.yaml new file mode 100644 index 000000000..fa200f7f6 --- /dev/null +++ b/releasenotes/notes/saturation-greedy-color-109d40f189590d3a.yaml @@ -0,0 +1,59 @@ +--- +features: + - Added a new class :class:`~.ColoringStrategy` used to specify the strategy used by + the greedy node and edge coloring algorithms. + The ``Degree`` strategy colors the nodes with higher degree first. + The ``Saturation`` strategy dynamically chooses the vertex that has the largest + number of different colors already assigned to its neighbors, and, in case of a tie, + the vertex that has the largest number of uncolored neighbors. + The ``IndependentSet`` strategy finds independent subsets of the graph, and assigns + a different color to each of these subsets. + - | + The rustworkx-core ``coloring`` module has 2 new functions, + ``greedy_node_color_with_coloring_strategy`` and + ``greedy_edge_color_with_coloring_strategy``. + These functions color respectively the nodes or the edges of the graph + using the specified coloring strategy and handling the preset colors + when provided. + - | + Added a new keyword argument, ``strategy``, to :func:`.graph_greedy_color` + and to :func:`.graph_greedy_edge_color` to specify the greedy coloring strategy. + + For example: + + .. jupyter-execute:: + + import rustworkx as rx + from rustworkx.visualization import mpl_draw + + graph = rx.generators.generalized_petersen_graph(5, 2) + + coloring = rx.graph_greedy_color(graph, strategy=rx.ColoringStrategy.Saturation) + colors = [coloring[node] for node in graph.node_indices()] + + layout = rx.shell_layout(graph, nlist=[[0, 1, 2, 3, 4],[6, 7, 8, 9, 5]]) + mpl_draw(graph, node_color=colors, pos=layout) + - | + Added a new keyword argument, ``preset_color_fn``, to :func:`.graph_greedy_edge_color` + which is used to provide preset colors for specific edges when computing the graph + coloring. You can optionally pass a callable to that argument which will + be passed edge index from the graph and is either expected to return an + integer color to use for that edge, or `None` to indicate there is no + preset color for that edge. For example: + + .. jupyter-execute:: + + import rustworkx as rx + from rustworkx.visualization import mpl_draw + + graph = rx.generators.generalized_petersen_graph(5, 2) + + def preset_colors(edge_index): + if edge_index == 0: + return 3 + + coloring = rx.graph_greedy_edge_color(graph, preset_color_fn=preset_colors) + colors = [coloring[edge] for edge in graph.edge_indices()] + + layout = rx.shell_layout(graph, nlist=[[0, 1, 2, 3, 4], [6, 7, 8, 9, 5]]) + mpl_draw(graph, edge_color=colors, pos=layout) diff --git a/rustworkx-core/src/coloring.rs b/rustworkx-core/src/coloring.rs index e0ca02ef4..730119edc 100644 --- a/rustworkx-core/src/coloring.rs +++ b/rustworkx-core/src/coloring.rs @@ -10,6 +10,9 @@ // License for the specific language governing permissions and limitations // under the License. +use ahash::RandomState; +use priority_queue::PriorityQueue; +use std::cmp::Ordering; use std::cmp::Reverse; use std::convert::Infallible; use std::hash::Hash; @@ -18,6 +21,7 @@ use crate::dictmap::*; use crate::line_graph::line_graph; use hashbrown::{HashMap, HashSet}; +use indexmap::IndexSet; use petgraph::graph::NodeIndex; use petgraph::visit::{ EdgeCount, EdgeIndexable, EdgeRef, GraphBase, GraphProp, IntoEdges, IntoNeighborsDirected, @@ -97,7 +101,36 @@ where Some(colors) } +/// coloring strategies for greedy node- and edge- coloring algorithms +#[derive(Clone, PartialEq)] +pub enum ColoringStrategy { + Degree, + Saturation, + IndependentSet, +} + fn inner_greedy_node_color( + graph: G, + preset_color_fn: F, + strategy: ColoringStrategy, +) -> Result, E> +where + G: NodeCount + IntoNodeIdentifiers + IntoEdges, + G::NodeId: Hash + Eq + Send + Sync, + F: FnMut(G::NodeId) -> Result, E>, +{ + match strategy { + ColoringStrategy::Degree => inner_greedy_node_color_strategy_degree(graph, preset_color_fn), + ColoringStrategy::Saturation => { + inner_greedy_node_color_strategy_saturation(graph, preset_color_fn) + } + ColoringStrategy::IndependentSet => { + inner_greedy_node_color_strategy_independent_set(graph, preset_color_fn) + } + } +} + +fn inner_greedy_node_color_strategy_degree( graph: G, mut preset_color_fn: F, ) -> Result, E> @@ -141,6 +174,163 @@ where Ok(colors) } +/// Data associated to nodes for the greedy coloring algorithm +/// using the "saturation first" strategy: always picking the node that +/// has the largest number of different colors already assigned to its +/// neighbors, and, in case of a tie, the node that has the largest number +/// of uncolored neighbors. +#[derive(Clone, Eq, PartialEq)] +struct SaturationStrategyData { + // degree of a node: number of neighbors without color + degree: usize, + // saturation degree of a node: number of colors used by neighbors + saturation: usize, +} + +impl Ord for SaturationStrategyData { + fn cmp(&self, other: &SaturationStrategyData) -> Ordering { + self.saturation + .cmp(&other.saturation) + .then_with(|| self.degree.cmp(&other.degree)) + } +} + +impl PartialOrd for SaturationStrategyData { + fn partial_cmp(&self, other: &SaturationStrategyData) -> Option { + Some(self.cmp(other)) + } +} + +fn inner_greedy_node_color_strategy_saturation( + graph: G, + mut preset_color_fn: F, +) -> Result, E> +where + G: NodeCount + IntoNodeIdentifiers + IntoEdges, + G::NodeId: Hash + Eq + Send + Sync, + F: FnMut(G::NodeId) -> Result, E>, +{ + let mut colors: DictMap = DictMap::with_capacity(graph.node_count()); + let mut nbd_colors: HashMap> = graph + .node_identifiers() + .map(|k| (k, HashSet::new())) + .collect(); + + let mut pq: PriorityQueue = + PriorityQueue::with_capacity(graph.node_count()); + + // Handle preset nodes + for k in graph.node_identifiers() { + if let Some(color) = preset_color_fn(k)? { + colors.insert(k, color); + for v in graph.neighbors(k) { + nbd_colors.get_mut(&v).unwrap().insert(color); + } + } + } + + // Add non-preset nodes to priority queue + for k in graph.node_identifiers() { + if colors.get(&k).is_none() { + let degree = graph + .neighbors(k) + .filter(|v| colors.get(v).is_none()) + .count(); + let saturation = nbd_colors.get(&k).unwrap().len(); + pq.push(k, SaturationStrategyData { degree, saturation }); + } + } + + // Greedily process nodes + while let Some((k, _)) = pq.pop() { + let neighbor_colors = nbd_colors.get(&k).unwrap(); + let mut current_color: usize = 0; + while neighbor_colors.contains(¤t_color) { + current_color += 1; + } + + colors.insert(k, current_color); + for v in graph.neighbors(k) { + if colors.get(&v).is_none() { + nbd_colors.get_mut(&v).unwrap().insert(current_color); + let (_, vdata) = pq.get(&v).unwrap(); + + pq.push( + v, + SaturationStrategyData { + degree: vdata.degree - 1, + saturation: nbd_colors.get(&v).unwrap().len(), + }, + ); + } + } + } + + Ok(colors) +} + +fn inner_greedy_node_color_strategy_independent_set( + graph: G, + mut preset_color_fn: F, +) -> Result, E> +where + G: NodeCount + IntoNodeIdentifiers + IntoEdges, + G::NodeId: Hash + Eq + Send + Sync, + F: FnMut(G::NodeId) -> Result, E>, +{ + let mut colors: DictMap = DictMap::with_capacity(graph.node_count()); + + let mut preset: HashSet = HashSet::new(); + let mut unprocessed: IndexSet = + IndexSet::with_hasher(RandomState::default()); + + // Handle preset nodes + for k in graph.node_identifiers() { + if let Some(color) = preset_color_fn(k)? { + colors.insert(k, color); + preset.insert(k); + } else { + unprocessed.insert(k); + } + } + + let mut current_color = 0; + while !unprocessed.is_empty() { + let mut remaining: IndexSet = + IndexSet::with_hasher(RandomState::default()); + + // Remove neighbors of preset nodes with the given color + for k in &preset { + if colors.get(k) == Some(¤t_color) { + for u in graph.neighbors(*k) { + if unprocessed.swap_take(&u).is_some() { + remaining.insert(u); + } + } + } + } + + // Greedily extract maximal independent set + while !unprocessed.is_empty() { + // Greedily take any node + // Possible optimization is to choose node with smallest degree among unprocessed + let k = *unprocessed.iter().next().unwrap(); + colors.insert(k, current_color); + unprocessed.swap_remove(&k); + for u in graph.neighbors(k) { + if unprocessed.swap_take(&u).is_some() { + remaining.insert(u); + } + } + } + + unprocessed = remaining; + current_color += 1; + } + + Ok(colors) +} + /// Color a graph using a greedy graph coloring algorithm. /// /// This function uses a `largest-first` strategy as described in: @@ -153,6 +343,10 @@ where /// The coloring problem is NP-hard and this is a heuristic algorithm /// which may not return an optimal solution. /// +/// # Note: +/// Please consider using ``greedy_node_color_with_coloring_strategy``, which is +/// a more general version of this function. +/// /// Arguments: /// /// * `graph` - The graph object to run the algorithm on @@ -180,7 +374,12 @@ where G: NodeCount + IntoNodeIdentifiers + IntoEdges, G::NodeId: Hash + Eq + Send + Sync, { - inner_greedy_node_color(graph, |_| Ok::, Infallible>(None)).unwrap() + inner_greedy_node_color( + graph, + |_| Ok::, Infallible>(None), + ColoringStrategy::Degree, + ) + .unwrap() } /// Color a graph using a greedy graph coloring algorithm with preset colors @@ -195,6 +394,10 @@ where /// The coloring problem is NP-hard and this is a heuristic algorithm /// which may not return an optimal solution. /// +/// # Note: +/// Please consider using ``greedy_node_color_with_coloring_strategy``, which is +/// a more general version of this function. +/// /// Arguments: /// /// * `graph` - The graph object to run the algorithm on @@ -238,7 +441,75 @@ where G::NodeId: Hash + Eq + Send + Sync, F: FnMut(G::NodeId) -> Result, E>, { - inner_greedy_node_color(graph, preset_color_fn) + inner_greedy_node_color(graph, preset_color_fn, ColoringStrategy::Degree) +} + +/// Color a graph using a greedy graph coloring algorithm with preset colors. +/// +/// This function uses one of several greedy strategies described in: +/// +/// Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs, +/// Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4. +/// +/// The `Degree` (aka `largest-first`) strategy colors the nodes with higher degree +/// first. The `Saturation` (aka `DSATUR` and `SLF`) strategy dynamically +/// chooses the vertex that has the largest number of different colors already +/// assigned to its neighbors, and, in case of a tie, the vertex that has the +/// largest number of uncolored neighbors. The `IndependentSet` strategy finds +/// independent subsets of the graph and assigns a different color to each of these +/// subsets. +/// +/// to color the nodes with higher degree first. +/// +/// The coloring problem is NP-hard and this is a heuristic algorithm +/// which may not return an optimal solution. +/// +/// Arguments: +/// +/// * `graph` - The graph object to run the algorithm on. +/// * `preset_color_fn` - A callback function that will receive the node identifier +/// for each node in the graph and is expected to return an `Option` +/// (wrapped in a `Result`) that is `None` if the node has no preset and +/// the usize represents the preset color. +/// * `strategy` - The greedy strategy used by the algorithm. +/// +/// # Example +/// ```rust +/// +/// use petgraph::graph::Graph; +/// use petgraph::graph::NodeIndex; +/// use petgraph::Undirected; +/// use rustworkx_core::dictmap::*; +/// use std::convert::Infallible; +/// use rustworkx_core::coloring::{greedy_node_color_with_coloring_strategy, ColoringStrategy}; +/// +/// let preset_color_fn = |node_idx: NodeIndex| -> Result, Infallible> { +/// if node_idx.index() == 0 { +/// Ok(Some(1)) +/// } else { +/// Ok(None) +/// } +/// }; +/// +/// let g = Graph::<(), (), Undirected>::from_edges(&[(0, 1), (0, 2)]); +/// let colors = greedy_node_color_with_coloring_strategy(&g, preset_color_fn, ColoringStrategy::Degree).unwrap(); +/// let mut expected_colors = DictMap::new(); +/// expected_colors.insert(NodeIndex::new(0), 1); +/// expected_colors.insert(NodeIndex::new(1), 0); +/// expected_colors.insert(NodeIndex::new(2), 0); +/// assert_eq!(colors, expected_colors); +/// ``` +pub fn greedy_node_color_with_coloring_strategy( + graph: G, + preset_color_fn: F, + strategy: ColoringStrategy, +) -> Result, E> +where + G: NodeCount + IntoNodeIdentifiers + IntoEdges, + G::NodeId: Hash + Eq + Send + Sync, + F: FnMut(G::NodeId) -> Result, E>, +{ + inner_greedy_node_color(graph, preset_color_fn, strategy) } /// Color edges of a graph using a greedy approach. @@ -248,6 +519,9 @@ where /// The coloring problem is NP-hard and this is a heuristic algorithm /// which may not return an optimal solution. /// +/// # Note: +/// Please consider using ``greedy_edge_color_with_coloring_strategy``, which is +/// a more general version of this function. /// Arguments: /// /// * `graph` - The graph object to run the algorithm on @@ -295,6 +569,78 @@ where edge_colors } +/// Color edges of a graph using a greedy graph coloring algorithm with preset +/// colors. +/// +/// This function works by greedily coloring the line graph of the given graph. +/// +/// The coloring problem is NP-hard and this is a heuristic algorithm +/// which may not return an optimal solution. +/// +/// Arguments: +/// +/// * `graph` - The graph object to run the algorithm on. +/// * `preset_color_fn` - A callback function that will receive the edge identifier +/// for each edge in the graph and is expected to return an `Option` +/// (wrapped in a `Result`) that is `None` if the edge has no preset and +/// the usize represents the preset color. +/// * `strategy` - The greedy strategy used by the algorithm. +/// +/// # Example +/// ```rust +/// +/// use petgraph::graph::Graph; +/// use petgraph::graph::EdgeIndex; +/// use petgraph::Undirected; +/// use rustworkx_core::dictmap::*; +/// use rustworkx_core::coloring::{greedy_edge_color_with_coloring_strategy, ColoringStrategy}; +/// use std::convert::Infallible; +/// +/// let g = Graph::<(), (), Undirected>::from_edges(&[(0, 1), (1, 2), (0, 2), (2, 3)]); +/// let preset_color_fn = |_| Ok::, Infallible>(None); +/// let colors = greedy_edge_color_with_coloring_strategy(&g, preset_color_fn, ColoringStrategy::Degree).unwrap(); +/// let mut expected_colors = DictMap::new(); +/// expected_colors.insert(EdgeIndex::new(0), 2); +/// expected_colors.insert(EdgeIndex::new(1), 0); +/// expected_colors.insert(EdgeIndex::new(2), 1); +/// expected_colors.insert(EdgeIndex::new(3), 2); +/// assert_eq!(colors, expected_colors); +/// ``` +/// +pub fn greedy_edge_color_with_coloring_strategy( + graph: G, + preset_color_fn: F, + strategy: ColoringStrategy, +) -> Result, E> +where + G: EdgeCount + IntoNodeIdentifiers + IntoEdges, + G::EdgeId: Hash + Eq, + F: Fn(G::EdgeId) -> Result, E>, +{ + let (new_graph, edge_to_node_map): ( + petgraph::graph::UnGraph<(), ()>, + HashMap, + ) = line_graph(&graph, || (), || ()); + + let node_to_edge_map: HashMap<&NodeIndex, &G::EdgeId> = + edge_to_node_map.iter().map(|(k, v)| (v, k)).collect(); + let new_graph_preset_color_fn = + |x: NodeIndex| preset_color_fn(**node_to_edge_map.get(&x).unwrap()); + + let colors = inner_greedy_node_color(&new_graph, new_graph_preset_color_fn, strategy)?; + + let mut edge_colors: DictMap = DictMap::with_capacity(graph.edge_count()); + + for edge in graph.edge_references() { + let edge_index = edge.id(); + let node_index = edge_to_node_map.get(&edge_index).unwrap(); + let edge_color = colors.get(node_index).unwrap(); + edge_colors.insert(edge_index, *edge_color); + } + + Ok(edge_colors) +} + struct MisraGries { // The input graph graph: G, @@ -559,17 +905,68 @@ where #[cfg(test)] mod test_node_coloring { - - use crate::coloring::greedy_node_color; - use crate::coloring::two_color; + use crate::coloring::{ + greedy_node_color, greedy_node_color_with_coloring_strategy, two_color, ColoringStrategy, + }; use crate::dictmap::*; - use crate::petgraph::prelude::*; + use crate::generators::{complete_graph, cycle_graph, heavy_hex_graph, path_graph}; + use std::convert::Infallible; + use std::hash::Hash; use crate::petgraph::graph::NodeIndex; + use crate::petgraph::prelude::*; use crate::petgraph::Undirected; + use petgraph::visit::{IntoEdgeReferences, IntoNodeIdentifiers}; + + /// Helper function to check validity of node coloring + fn check_node_colors(graph: G, colors: &DictMap) + where + G: IntoNodeIdentifiers + IntoEdgeReferences, + G::NodeId: Hash + Eq + Send + Sync, + { + // Check that every node has valid color + for k in graph.node_identifiers() { + if !colors.contains_key(&k) { + panic!("Problem: some nodes have no color assigned."); + } else { + println!("Valid color: ok"); + } + } + + // Check that nodes connected by an edge have different colors + for e in graph.edge_references() { + if colors.get(&e.source()) == colors.get(&e.target()) { + panic!("Problem: same color for connected nodes."); + } else { + println!("Connected nodes: ok"); + } + } + } + + /// Helper function to check validity of node coloring with preset colors + fn check_preset_colors( + graph: G, + colors: &DictMap, + mut preset_color_fn: F, + ) where + G: IntoNodeIdentifiers + IntoEdgeReferences, + G::NodeId: Hash + Eq + Send + Sync, + F: FnMut(G::NodeId) -> Result, E>, + { + // Check preset values + for k in graph.node_identifiers() { + if let Ok(Some(color)) = preset_color_fn(k) { + if *colors.get(&k).unwrap() != color { + panic!("Problem: colors are different from preset vales."); + } else { + println!("Preset values: ok"); + } + } + } + } #[test] - fn test_greedy_node_color_empty_graph() { + fn test_legacy_greedy_node_color_empty_graph() { // Empty graph let graph = Graph::<(), (), Undirected>::new_undirected(); let colors = greedy_node_color(&graph); @@ -578,7 +975,7 @@ mod test_node_coloring { } #[test] - fn test_greedy_node_color_simple_graph() { + fn test_legacy_greedy_node_color_simple_graph() { // Simple graph let graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (0, 2)]); let colors = greedy_node_color(&graph); @@ -593,7 +990,7 @@ mod test_node_coloring { } #[test] - fn test_greedy_node_color_simple_graph_large_degree() { + fn test_legacy_greedy_node_color_simple_graph_large_degree() { // Graph with multiple edges let graph = Graph::<(), (), Undirected>::from_edges([ (0, 1), @@ -614,6 +1011,226 @@ mod test_node_coloring { assert_eq!(colors, expected_colors); } + #[test] + fn test_greedy_node_color_empty_graph() { + // Empty graph + let graph = Graph::<(), (), Undirected>::new_undirected(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + for strategy in vec![ + ColoringStrategy::Degree, + ColoringStrategy::Saturation, + ColoringStrategy::IndependentSet, + ] { + let colors = + greedy_node_color_with_coloring_strategy(&graph, preset_color_fn, strategy); + let expected_colors: DictMap = [].into_iter().collect(); + assert_eq!(colors, Ok(expected_colors)); + } + } + + #[test] + fn test_greedy_node_color_simple_graph() { + // Simple graph + let graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (0, 2)]); + let preset_color_fn = |_| Ok::, Infallible>(None); + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ); + let expected_colors: DictMap = [ + (NodeIndex::new(0), 0), + (NodeIndex::new(1), 1), + (NodeIndex::new(2), 1), + ] + .into_iter() + .collect(); + assert_eq!(colors, Ok(expected_colors)); + } + + #[test] + fn test_greedy_node_color_simple_graph_large_degree() { + // Graph with multiple edges + let graph = Graph::<(), (), Undirected>::from_edges([ + (0, 1), + (0, 2), + (0, 2), + (0, 2), + (0, 2), + (0, 2), + ]); + let preset_color_fn = |_| Ok::, Infallible>(None); + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ); + let expected_colors: DictMap = [ + (NodeIndex::new(0), 0), + (NodeIndex::new(1), 1), + (NodeIndex::new(2), 1), + ] + .into_iter() + .collect(); + assert_eq!(colors, Ok(expected_colors)); + } + + #[test] + fn test_greedy_node_color_saturation() { + // Simple graph + let graph = Graph::<(), (), Undirected>::from_edges([ + (0, 1), + (0, 2), + (0, 3), + (3, 4), + (4, 5), + (5, 6), + (5, 7), + ]); + + let preset_color_fn = |_| Ok::, Infallible>(None); + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Saturation, + ) + .unwrap(); + check_node_colors(&graph, &colors); + + let expected_colors: DictMap = [ + (NodeIndex::new(0), 0), + (NodeIndex::new(1), 1), + (NodeIndex::new(2), 1), + (NodeIndex::new(3), 1), + (NodeIndex::new(4), 0), + (NodeIndex::new(5), 1), + (NodeIndex::new(6), 0), + (NodeIndex::new(7), 0), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_node_color_saturation_and_preset() { + // Simple graph + let graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (0, 2), (2, 3), (2, 4)]); + + let preset_color_fn = |node_idx: NodeIndex| -> Result, Infallible> { + if node_idx.index() == 0 { + Ok(Some(1)) + } else { + Ok(None) + } + }; + + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Saturation, + ) + .unwrap(); + check_node_colors(&graph, &colors); + check_preset_colors(&graph, &colors, preset_color_fn); + + let expected_colors: DictMap = [ + (NodeIndex::new(0), 1), + (NodeIndex::new(1), 0), + (NodeIndex::new(2), 0), + (NodeIndex::new(3), 1), + (NodeIndex::new(4), 1), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_node_color_independent_set() { + // Simple graph + let graph = Graph::<(), (), Undirected>::from_edges([ + (0, 1), + (0, 2), + (0, 3), + (3, 4), + (4, 5), + (5, 6), + (5, 7), + ]); + + let preset_color_fn = |_| Ok::, Infallible>(None); + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::IndependentSet, + ) + .unwrap(); + check_node_colors(&graph, &colors); + + let expected_colors: DictMap = [ + (NodeIndex::new(0), 0), + (NodeIndex::new(1), 1), + (NodeIndex::new(2), 1), + (NodeIndex::new(3), 1), + (NodeIndex::new(4), 0), + (NodeIndex::new(5), 1), + (NodeIndex::new(6), 0), + (NodeIndex::new(7), 0), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_node_color_independent_set_and_preset() { + // Simple graph + let graph = Graph::<(), (), Undirected>::from_edges([ + (0, 1), + (0, 2), + (0, 3), + (3, 4), + (4, 5), + (5, 6), + (5, 7), + ]); + + let preset_color_fn = |node_idx: NodeIndex| -> Result, Infallible> { + if node_idx.index() == 0 { + Ok(Some(1)) + } else if node_idx.index() == 3 { + Ok(Some(0)) + } else { + Ok(None) + } + }; + + let colors = greedy_node_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::IndependentSet, + ) + .unwrap(); + check_node_colors(&graph, &colors); + check_preset_colors(&graph, &colors, preset_color_fn); + + let expected_colors: DictMap = [ + (NodeIndex::new(0), 1), + (NodeIndex::new(1), 0), + (NodeIndex::new(2), 0), + (NodeIndex::new(3), 0), + (NodeIndex::new(4), 1), + (NodeIndex::new(5), 2), + (NodeIndex::new(6), 0), + (NodeIndex::new(7), 0), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + #[test] fn test_two_color_directed() { let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 4)]; @@ -684,19 +1301,94 @@ mod test_node_coloring { expected_colors.insert(NodeIndex::new(6), 1); assert_eq!(coloring, expected_colors) } + + #[test] + fn test_path_graph() { + let graph: petgraph::graph::UnGraph<(), ()> = + path_graph(Some(7), None, || (), || (), false).unwrap(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + for strategy in vec![ + ColoringStrategy::Degree, + ColoringStrategy::Saturation, + ColoringStrategy::IndependentSet, + ] { + let colors = + greedy_node_color_with_coloring_strategy(&graph, preset_color_fn, strategy) + .unwrap(); + check_node_colors(&graph, &colors); + } + } + + #[test] + fn test_cycle_graph() { + let graph: petgraph::graph::UnGraph<(), ()> = + cycle_graph(Some(15), None, || (), || (), false).unwrap(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + for strategy in vec![ + ColoringStrategy::Degree, + ColoringStrategy::Saturation, + ColoringStrategy::IndependentSet, + ] { + let colors = + greedy_node_color_with_coloring_strategy(&graph, preset_color_fn, strategy) + .unwrap(); + check_node_colors(&graph, &colors); + } + } + + #[test] + fn test_heavy_hex_graph() { + let graph: petgraph::graph::UnGraph<(), ()> = + heavy_hex_graph(7, || (), || (), false).unwrap(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + for strategy in vec![ + ColoringStrategy::Degree, + ColoringStrategy::Saturation, + ColoringStrategy::IndependentSet, + ] { + let colors = + greedy_node_color_with_coloring_strategy(&graph, preset_color_fn, strategy) + .unwrap(); + check_node_colors(&graph, &colors); + } + } + + #[test] + fn test_complete_graph() { + let graph: petgraph::graph::UnGraph<(), ()> = + complete_graph(Some(10), None, || (), || ()).unwrap(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + for strategy in vec![ + ColoringStrategy::Degree, + ColoringStrategy::Saturation, + ColoringStrategy::IndependentSet, + ] { + let colors = + greedy_node_color_with_coloring_strategy(&graph, preset_color_fn, strategy) + .unwrap(); + check_node_colors(&graph, &colors); + } + } } #[cfg(test)] mod test_edge_coloring { - use crate::coloring::greedy_edge_color; + use crate::coloring::{ + greedy_edge_color, greedy_edge_color_with_coloring_strategy, ColoringStrategy, + }; use crate::dictmap::DictMap; use crate::petgraph::Graph; + use std::convert::Infallible; use petgraph::graph::{edge_index, EdgeIndex}; use petgraph::Undirected; #[test] - fn test_greedy_edge_color_empty_graph() { + fn test_legacy_greedy_edge_color_empty_graph() { // Empty graph let graph = Graph::<(), (), Undirected>::new_undirected(); let colors = greedy_edge_color(&graph); @@ -705,7 +1397,7 @@ mod test_edge_coloring { } #[test] - fn test_greedy_edge_color_simple_graph() { + fn test_legacy_greedy_edge_color_simple_graph() { // Graph with an edge removed let graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3)]); let colors = greedy_edge_color(&graph); @@ -720,7 +1412,7 @@ mod test_edge_coloring { } #[test] - fn test_greedy_edge_color_graph_with_removed_edges() { + fn test_legacy_greedy_edge_color_graph_with_removed_edges() { // Simple graph let mut graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0)]); graph.remove_edge(edge_index(1)); @@ -734,6 +1426,244 @@ mod test_edge_coloring { .collect(); assert_eq!(colors, expected_colors); } + + #[test] + fn test_greedy_edge_color_empty_graph() { + // Empty graph + let graph = Graph::<(), (), Undirected>::new_undirected(); + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ) + .unwrap(); + let expected_colors: DictMap = [].into_iter().collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_simple_graph() { + // Graph with an edge removed + let graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3)]); + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 1), + (EdgeIndex::new(1), 0), + (EdgeIndex::new(2), 1), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_graph_with_removed_edges() { + // Simple graph + let mut graph = Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0)]); + graph.remove_edge(edge_index(1)); + + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ) + .unwrap(); + + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 1), + (EdgeIndex::new(1), 0), + (EdgeIndex::new(2), 1), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_degree() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 1), + (EdgeIndex::new(1), 0), + (EdgeIndex::new(2), 1), + (EdgeIndex::new(3), 0), + (EdgeIndex::new(4), 2), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_degree_with_preset() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + + let preset_color_fn = |node_idx: EdgeIndex| -> Result, Infallible> { + if node_idx.index() == 1 { + Ok(Some(1)) + } else { + Ok(None) + } + }; + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Degree, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 0), + (EdgeIndex::new(1), 1), + (EdgeIndex::new(2), 0), + (EdgeIndex::new(3), 1), + (EdgeIndex::new(4), 2), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_saturation() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Saturation, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 1), + (EdgeIndex::new(1), 0), + (EdgeIndex::new(2), 1), + (EdgeIndex::new(3), 0), + (EdgeIndex::new(4), 2), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_saturation_with_preset() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + + let preset_color_fn = |node_idx: EdgeIndex| -> Result, Infallible> { + if node_idx.index() == 1 { + Ok(Some(1)) + } else if node_idx.index() == 4 { + Ok(Some(0)) + } else { + Ok(None) + } + }; + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::Saturation, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 0), + (EdgeIndex::new(1), 1), + (EdgeIndex::new(2), 2), + (EdgeIndex::new(3), 1), + (EdgeIndex::new(4), 0), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_independent_set() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + let preset_color_fn = |_| Ok::, Infallible>(None); + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::IndependentSet, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 0), + (EdgeIndex::new(1), 1), + (EdgeIndex::new(2), 2), + (EdgeIndex::new(3), 1), + (EdgeIndex::new(4), 0), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } + + #[test] + fn test_greedy_edge_color_independent_set_with_preset() { + // Simple graph + let graph = + Graph::<(), (), Undirected>::from_edges([(0, 1), (1, 2), (2, 3), (3, 0), (2, 4)]); + + let preset_color_fn = |node_idx: EdgeIndex| -> Result, Infallible> { + if node_idx.index() == 1 { + Ok(Some(0)) + } else if node_idx.index() == 4 { + Ok(Some(2)) + } else { + Ok(None) + } + }; + + let colors = greedy_edge_color_with_coloring_strategy( + &graph, + preset_color_fn, + ColoringStrategy::IndependentSet, + ) + .unwrap(); + let expected_colors: DictMap = [ + (EdgeIndex::new(0), 1), + (EdgeIndex::new(1), 0), + (EdgeIndex::new(2), 1), + (EdgeIndex::new(3), 0), + (EdgeIndex::new(4), 2), + ] + .into_iter() + .collect(); + assert_eq!(colors, expected_colors); + } } #[cfg(test)] diff --git a/rustworkx/__init__.pyi b/rustworkx/__init__.pyi index f499b9813..11edc5922 100644 --- a/rustworkx/__init__.pyi +++ b/rustworkx/__init__.pyi @@ -33,6 +33,7 @@ from .rustworkx import JSONSerializationError as JSONSerializationError from .rustworkx import FailedToConverge as FailedToConverge from .rustworkx import InvalidMapping as InvalidMapping from .rustworkx import GraphNotBipartite as GraphNotBipartite +from .rustworkx import ColoringStrategy as ColoringStrategy from .rustworkx import digraph_cartesian_product as digraph_cartesian_product from .rustworkx import graph_cartesian_product as graph_cartesian_product diff --git a/rustworkx/rustworkx.pyi b/rustworkx/rustworkx.pyi index 522962994..47c8e3673 100644 --- a/rustworkx/rustworkx.pyi +++ b/rustworkx/rustworkx.pyi @@ -50,6 +50,12 @@ class FailedToConverge(Exception): ... class InvalidMapping(Exception): ... class GraphNotBipartite(Exception): ... +@final +class ColoringStrategy: + Degree: Any + Saturation: Any + IndependentSet: Any + # Cartesian product def digraph_cartesian_product( @@ -139,9 +145,17 @@ def graph_katz_centrality( # Coloring def graph_greedy_color( - graph: PyGraph, /, preset_color_fn: Callable[[int], int | None] | None = ... + graph: PyGraph, + /, + preset_color_fn: Callable[[int], int | None] | None = ..., + strategy: int = ..., +) -> dict[int, int]: ... +def graph_greedy_edge_color( + graph: PyGraph, + /, + preset_color_fn: Callable[[int], int | None] | None = ..., + strategy: int = ..., ) -> dict[int, int]: ... -def graph_greedy_edge_color(graph: PyGraph, /) -> dict[int, int]: ... def graph_is_bipartite(graph: PyGraph) -> bool: ... def digraph_is_bipartite(graph: PyDiGraph) -> bool: ... def graph_two_color(graph: PyGraph) -> dict[int, int]: ... diff --git a/src/coloring.rs b/src/coloring.rs index b544ffa5f..cf844b1a2 100644 --- a/src/coloring.rs +++ b/src/coloring.rs @@ -11,29 +11,62 @@ // under the License. use crate::GraphNotBipartite; -use crate::{digraph, graph, NodeIndex}; +use crate::{digraph, graph, EdgeIndex, NodeIndex}; use pyo3::prelude::*; use pyo3::types::PyDict; use pyo3::Python; +use std::convert::Infallible; + use rustworkx_core::bipartite_coloring::bipartite_edge_color; use rustworkx_core::coloring::{ - greedy_edge_color, greedy_node_color, greedy_node_color_with_preset_colors, + greedy_edge_color_with_coloring_strategy, greedy_node_color_with_coloring_strategy, misra_gries_edge_color, two_color, }; +pub use rustworkx_core::coloring::ColoringStrategy as ColoringStrategyCore; + +/// Greedy coloring strategies available for `graph_greedy_color` +/// +/// .. list-table:: Strategy description +/// :header-rows: 1 +/// +/// * - Strategy +/// - Reference +/// * - Degree +/// - `Largest-first` strategy in [1] (section 1.2.2.2) +/// * - Saturation +/// - `DSATUR` strategy in [1] (section 1.2.2.8) +/// * - IndependentSet +/// - `GIS` strategy in [1] (section 1.2.2.9) +/// +/// [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs, Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4. +#[pyclass(module = "rustworkx")] +#[derive(Clone, PartialEq)] +pub enum ColoringStrategy { + Degree, + Saturation, + IndependentSet, +} + /// Color a :class:`~.PyGraph` object using a greedy graph coloring algorithm. /// -/// This function uses a `largest-first` strategy as described in [1]_ and colors -/// the nodes with higher degree first. +/// This function uses one of several greedy strategies described in: [1]_. +/// The `Degree` (aka `largest-first`) strategy colors the nodes with higher degree +/// first. The `Saturation` (aka `DSATUR` and `SLF`) strategy dynamically +/// chooses the vertex that has the largest number of different colors already +/// assigned to its neighbors, and, in case of a tie, the vertex that has the +/// largest number of uncolored neighbors. The `IndependentSet` strategy finds +/// independent subsets of the graph and assigns a different color to each of these +/// subsets. /// /// .. note:: /// /// The coloring problem is NP-hard and this is a heuristic algorithm which /// may not return an optimal solution. /// -/// :param PyGraph: The input PyGraph object to color +/// :param PyGraph: The input PyGraph object to color. /// :param preset_color_fn: An optional callback function that is used to manually /// specify a color to use for particular nodes in the graph. If specified /// this takes a callable that will be passed a node index and is expected to @@ -41,6 +74,9 @@ use rustworkx_core::coloring::{ /// is no preset. Note if you do use a callable there is no validation that /// the preset values are valid colors. You can generate an invalid coloring /// if you the specified function returned invalid colors for any nodes. +/// :param strategy: The strategy used by the algorithm. When the +/// strategy is not explicitly specified, the `Degree` strategy is used by +/// default. /// /// :returns: A dictionary where keys are node indices and the value is /// the color @@ -63,12 +99,20 @@ use rustworkx_core::coloring::{ /// .. [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs, /// Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4. #[pyfunction] -#[pyo3(text_signature = "(graph, /, preset_color_fn=None)")] +#[pyo3(text_signature = "(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree)")] +#[pyo3(signature=(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree))] pub fn graph_greedy_color( py: Python, graph: &graph::PyGraph, preset_color_fn: Option, + strategy: ColoringStrategy, ) -> PyResult { + let inner_strategy = match strategy { + ColoringStrategy::Saturation => ColoringStrategyCore::Saturation, + ColoringStrategy::Degree => ColoringStrategyCore::Degree, + ColoringStrategy::IndependentSet => ColoringStrategyCore::IndependentSet, + }; + let colors = match preset_color_fn { Some(preset_color_fn) => { let callback = |node_idx: NodeIndex| -> PyResult> { @@ -76,9 +120,12 @@ pub fn graph_greedy_color( .call1(py, (node_idx.index(),)) .map(|x| x.extract(py).ok()) }; - greedy_node_color_with_preset_colors(&graph.graph, callback)? + greedy_node_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)? + } + None => { + let callback = |_: NodeIndex| -> Result, Infallible> { Ok(None) }; + greedy_node_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)? } - None => greedy_node_color(&graph.graph), }; let out_dict = PyDict::new_bound(py); for (node, color) in colors { @@ -91,7 +138,26 @@ pub fn graph_greedy_color( /// /// This function works by greedily coloring the line graph of the given graph. /// -/// :param PyGraph: The input PyGraph object to edge-color +/// This function uses one of several greedy strategies described in: [1]_. +/// The `Degree` (aka `largest-first`) strategy colors the nodes with higher degree +/// first. The `Saturation` (aka `DSATUR` and `SLF`) strategy dynamically +/// chooses the vertex that has the largest number of different colors already +/// assigned to its neighbors, and, in case of a tie, the vertex that has the +/// largest number of uncolored neighbors. The `IndependentSet` strategy finds +/// independent subsets of the graph and assigns a different color to each of these +/// subsets. +/// +/// :param PyGraph: The input PyGraph object to edge-color. +/// :param preset_color_fn: An optional callback function that is used to manually +/// specify a color to use for particular edges in the graph. If specified +/// this takes a callable that will be passed an edge index and is expected to +/// either return an integer representing a color or ``None`` to indicate there +/// is no preset. Note if you do use a callable there is no validation that +/// the preset values are valid colors. You can generate an invalid coloring +/// if you the specified function returned invalid colors for any edges. +/// :param strategy: The greedy strategy used by the algorithm. When the +/// strategy is not explicitly specified, the `Degree` strategy is used by +/// default. /// /// :returns: A dictionary where keys are edge indices and the value is the color /// :rtype: dict @@ -104,10 +170,39 @@ pub fn graph_greedy_color( /// edge_colors = rx.graph_greedy_edge_color(graph) /// assert edge_colors == {0: 0, 1: 1, 2: 0, 3: 1, 4: 0, 5: 1, 6: 2} /// +/// +/// .. [1] Adrian Kosowski, and Krzysztof Manuszewski, Classical Coloring of Graphs, +/// Graph Colorings, 2-19, 2004. ISBN 0-8218-3458-4. #[pyfunction] -#[pyo3(text_signature = "(graph, /)")] -pub fn graph_greedy_edge_color(py: Python, graph: &graph::PyGraph) -> PyResult { - let colors = greedy_edge_color(&graph.graph); +#[pyo3(text_signature = "(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree)")] +#[pyo3(signature=(graph, /, preset_color_fn=None, strategy=ColoringStrategy::Degree))] +pub fn graph_greedy_edge_color( + py: Python, + graph: &graph::PyGraph, + preset_color_fn: Option, + strategy: ColoringStrategy, +) -> PyResult { + let inner_strategy = match strategy { + ColoringStrategy::Saturation => ColoringStrategyCore::Saturation, + ColoringStrategy::Degree => ColoringStrategyCore::Degree, + ColoringStrategy::IndependentSet => ColoringStrategyCore::IndependentSet, + }; + + let colors = match preset_color_fn { + Some(preset_color_fn) => { + let callback = |edge_idx: EdgeIndex| -> PyResult> { + preset_color_fn + .call1(py, (edge_idx.index(),)) + .map(|x| x.extract(py).ok()) + }; + greedy_edge_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)? + } + None => { + let callback = |_: EdgeIndex| -> Result, Infallible> { Ok(None) }; + greedy_edge_color_with_coloring_strategy(&graph.graph, callback, inner_strategy)? + } + }; + let out_dict = PyDict::new_bound(py); for (node, color) in colors { out_dict.set_item(node.index(), color)?; diff --git a/src/lib.rs b/src/lib.rs index 164b713c5..cce7c9175 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -610,6 +610,7 @@ fn rustworkx(py: Python<'_>, m: &Bound) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_wrapped(wrap_pymodule!(generators::generators))?; Ok(()) } diff --git a/tests/graph/test_coloring.py b/tests/graph/test_coloring.py index 1cba11f9a..1899be756 100644 --- a/tests/graph/test_coloring.py +++ b/tests/graph/test_coloring.py @@ -87,6 +87,33 @@ def preset(node_idx): with self.assertRaises(OverflowError): rustworkx.graph_greedy_color(graph, preset) + def test_greedy_strategies(self): + graph = rustworkx.PyGraph() + [a, b, c, d, e, f, g, h] = graph.add_nodes_from(["a", "b", "c", "d", "e", "f", "g", "h"]) + graph.add_edges_from( + [(a, b, 1), (a, c, 1), (a, d, 1), (d, e, 1), (e, f, 1), (f, g, 1), (f, h, 1)] + ) + + with self.subTest(): + res = rustworkx.graph_greedy_color(graph) + self.assertEqual({a: 0, b: 1, c: 1, d: 1, e: 2, f: 0, g: 1, h: 1}, res) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Degree): + res = rustworkx.graph_greedy_color(graph, strategy=rustworkx.ColoringStrategy.Degree) + self.assertEqual({a: 0, b: 1, c: 1, d: 1, e: 2, f: 0, g: 1, h: 1}, res) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Saturation): + res = rustworkx.graph_greedy_color( + graph, strategy=rustworkx.ColoringStrategy.Saturation + ) + self.assertEqual({a: 0, b: 1, c: 1, d: 1, e: 0, f: 1, g: 0, h: 0}, res) + + with self.subTest(strategy=rustworkx.ColoringStrategy.IndependentSet): + res = rustworkx.graph_greedy_color( + graph, strategy=rustworkx.ColoringStrategy.IndependentSet + ) + self.assertEqual({a: 0, b: 1, c: 1, d: 1, e: 0, f: 1, g: 0, h: 0}, res) + class TestGraphEdgeColoring(unittest.TestCase): def test_graph(self): @@ -149,6 +176,67 @@ def test_cycle_graph(self): edge_colors = rustworkx.graph_greedy_edge_color(graph) self.assertEqual({0: 0, 1: 1, 2: 0, 3: 1, 4: 0, 5: 1, 6: 2}, edge_colors) + def test_greedy_strategies(self): + graph = rustworkx.generators.complete_graph(4) + + with self.subTest(): + edge_colors = rustworkx.graph_greedy_edge_color(graph) + self.assertEqual({0: 0, 1: 1, 2: 2, 3: 2, 4: 1, 5: 0}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Degree): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, strategy=rustworkx.ColoringStrategy.Degree + ) + self.assertEqual({0: 0, 1: 1, 2: 2, 3: 2, 4: 1, 5: 0}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Saturation): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, strategy=rustworkx.ColoringStrategy.Saturation + ) + self.assertEqual({0: 0, 1: 2, 2: 1, 3: 1, 4: 2, 5: 0}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.IndependentSet): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, strategy=rustworkx.ColoringStrategy.IndependentSet + ) + self.assertEqual({0: 0, 1: 2, 2: 1, 3: 1, 4: 2, 5: 0}, edge_colors) + + def test_greedy_strategies_with_preset(self): + def preset(edge_idx): + if edge_idx == 0: + return 1 + elif edge_idx == 3: + return 0 + else: + return None + + graph = rustworkx.generators.complete_graph(4) + + with self.subTest(): + edge_colors = rustworkx.graph_greedy_edge_color(graph, preset_color_fn=preset) + print(f"{edge_colors = }") + self.assertEqual({0: 1, 1: 2, 2: 0, 3: 0, 4: 2, 5: 1}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Degree): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, preset_color_fn=preset, strategy=rustworkx.ColoringStrategy.Degree + ) + self.assertEqual({0: 1, 1: 2, 2: 0, 3: 0, 4: 2, 5: 1}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.Saturation): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, preset_color_fn=preset, strategy=rustworkx.ColoringStrategy.Saturation + ) + self.assertEqual({0: 1, 1: 2, 2: 0, 3: 0, 4: 2, 5: 1}, edge_colors) + + with self.subTest(strategy=rustworkx.ColoringStrategy.IndependentSet): + edge_colors = rustworkx.graph_greedy_edge_color( + graph, + preset_color_fn=preset, + strategy=rustworkx.ColoringStrategy.IndependentSet, + ) + self.assertEqual({0: 1, 1: 2, 2: 0, 3: 0, 4: 2, 5: 1}, edge_colors) + class TestMisraGriesColoring(unittest.TestCase): def test_simple_graph(self):