Add two_color and is_bipartite (#1002)

* Add two_color and is_bipartite

This commit adds new functions to rustworkx, two_color() and
is_bipartite(), which are used to compute a two coloring for a graph and
then determine if a givn graph is bipartite. The two_color() function is
added to rustworkx-core as the python is_bipartite() function just wraps
it and converts the output to a bool if a two coloring is possible or
not.

This commit is based on top of #998 and will need to be rebased after
that merges.

* Remove special isolates handling

* Expand test coverage

releasenotes/notes/add-bipartite-9df3898a156e799c.yaml

@@ -0,0 +1,22 @@
+---
+features:
+  - |
+    Added a new function ``two_color`` to the rustworkx-core ``rustworkx_core::coloring``
+    module. This function is used to compute a two coloring of a graph and can
+    also be used to determine if a graph is bipartite as it returns ``None``
+    when a two coloring is not possible.
+  - |
+    Added a new function, :func:`~.two_color`, which is used to compute a
+    two coloring for a graph. For example:
+
+    .. jupyter-execute::
+
+      import rustworkx as rx
+      from rustworkx.visualization import mpl_draw
+
+      graph = rx.generators.heavy_square_graph(5)
+      colors = rx.two_color(graph)
+      mpl_draw(graph, node_color=list(colors.values()))
+  - |
+    Added a new function, :func:`~.is_bipartite` to determine whether a given
+    graph object is bipartite or not.

rustworkx-core/src/coloring.rs

@@ -17,9 +17,84 @@ use crate::dictmap::*;
 use crate::line_graph::line_graph;
 use hashbrown::{HashMap, HashSet};
 use petgraph::graph::NodeIndex;
-use petgraph::visit::{EdgeCount, EdgeRef, IntoEdges, IntoNodeIdentifiers, NodeCount};
+use petgraph::visit::{
+    EdgeCount, EdgeRef, GraphBase, GraphProp, IntoEdges, IntoNeighborsDirected,
+    IntoNodeIdentifiers, NodeCount, NodeIndexable,
+};
+use petgraph::{Incoming, Outgoing};
 use rayon::prelude::*;
 
+/// Compute a two-coloring of a graph
+///
+/// If a two coloring is not possible for the input graph (meaning it is not
+/// bipartite), `None` is returned.
+///
+/// Arguments:
+///
+/// * `graph` - The graph to find the coloring for
+///
+/// # Example
+///
+/// ```rust
+/// use rustworkx_core::petgraph::prelude::*;
+/// use rustworkx_core::coloring::two_color;
+/// use rustworkx_core::dictmap::*;
+///
+/// let edge_list = vec![
+///  (0, 1),
+///  (1, 2),
+///  (2, 3),
+///  (3, 4),
+/// ];
+///
+/// let graph = UnGraph::<i32, i32>::from_edges(&edge_list);
+/// let coloring = two_color(&graph).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), 1);
+/// expected_colors.insert(NodeIndex::new(3), 0);
+/// expected_colors.insert(NodeIndex::new(4), 1);
+/// assert_eq!(coloring, expected_colors)
+/// ```
+pub fn two_color<G>(graph: G) -> Option<DictMap<G::NodeId, u8>>
+where
+    G: NodeIndexable
+        + IntoNodeIdentifiers
+        + IntoNeighborsDirected
+        + GraphBase
+        + GraphProp
+        + NodeCount,
+    <G as GraphBase>::NodeId: std::cmp::Eq + Hash,
+{
+    let mut colors = DictMap::with_capacity(graph.node_count());
+    for node in graph.node_identifiers() {
+        if colors.contains_key(&node) {
+            continue;
+        }
+        let mut queue = vec![node];
+        colors.insert(node, 1);
+        while let Some(v) = queue.pop() {
+            let v_color: u8 = *colors.get(&v).unwrap();
+            let color: u8 = 1 - v_color;
+            for w in graph
+                .neighbors_directed(v, Outgoing)
+                .chain(graph.neighbors_directed(v, Incoming))
+            {
+                if let Some(color_w) = colors.get(&w) {
+                    if *color_w == v_color {
+                        return None;
+                    }
+                } else {
+                    colors.insert(w, color);
+                    queue.push(w);
+                }
+            }
+        }
+    }
+    Some(colors)
+}
+
 /// Color a graph using a greedy graph coloring algorithm.
 ///
 /// This function uses a `largest-first` strategy as described in:
@@ -150,11 +225,12 @@ where
 mod test_node_coloring {
 
     use crate::coloring::greedy_node_color;
-    use crate::dictmap::DictMap;
-    use crate::petgraph::Graph;
+    use crate::coloring::two_color;
+    use crate::dictmap::*;
+    use crate::petgraph::prelude::*;
 
-    use petgraph::graph::NodeIndex;
-    use petgraph::Undirected;
+    use crate::petgraph::graph::NodeIndex;
+    use crate::petgraph::Undirected;
 
     #[test]
     fn test_greedy_node_color_empty_graph() {
@@ -201,6 +277,77 @@ mod test_node_coloring {
         .collect();
         assert_eq!(colors, expected_colors);
     }
+
+    #[test]
+    fn test_two_color_directed() {
+        let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 4)];
+
+        let graph = DiGraph::<i32, i32>::from_edges(&edge_list);
+        let coloring = two_color(&graph).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), 1);
+        expected_colors.insert(NodeIndex::new(3), 0);
+        expected_colors.insert(NodeIndex::new(4), 1);
+        assert_eq!(coloring, expected_colors)
+    }
+
+    #[test]
+    fn test_two_color_directed_not_bipartite() {
+        let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 0), (3, 1)];
+
+        let graph = DiGraph::<i32, i32>::from_edges(&edge_list);
+        let coloring = two_color(&graph);
+        assert_eq!(None, coloring)
+    }
+
+    #[test]
+    fn test_two_color_undirected_not_bipartite() {
+        let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 0), (3, 1)];
+
+        let graph = UnGraph::<i32, i32>::from_edges(&edge_list);
+        let coloring = two_color(&graph);
+        assert_eq!(None, coloring)
+    }
+
+    #[test]
+    fn test_two_color_directed_with_isolates() {
+        let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 4)];
+
+        let mut graph = DiGraph::<i32, i32>::from_edges(&edge_list);
+        graph.add_node(10);
+        graph.add_node(11);
+        let coloring = two_color(&graph).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), 1);
+        expected_colors.insert(NodeIndex::new(3), 0);
+        expected_colors.insert(NodeIndex::new(4), 1);
+        expected_colors.insert(NodeIndex::new(5), 1);
+        expected_colors.insert(NodeIndex::new(6), 1);
+        assert_eq!(coloring, expected_colors)
+    }
+
+    #[test]
+    fn test_two_color_undirected_with_isolates() {
+        let edge_list = vec![(0, 1), (1, 2), (2, 3), (3, 4)];
+
+        let mut graph = UnGraph::<i32, i32>::from_edges(&edge_list);
+        graph.add_node(10);
+        graph.add_node(11);
+        let coloring = two_color(&graph).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), 1);
+        expected_colors.insert(NodeIndex::new(3), 0);
+        expected_colors.insert(NodeIndex::new(4), 1);
+        expected_colors.insert(NodeIndex::new(5), 1);
+        expected_colors.insert(NodeIndex::new(6), 1);
+        assert_eq!(coloring, expected_colors)
+    }
 }
 
 #[cfg(test)]

rustworkx/__init__.py

@@ -2071,3 +2071,34 @@ def isolates(graph):
 
 isolates.register(PyDiGraph, digraph_isolates)
 isolates.register(PyGraph, graph_isolates)
+
+
+@functools.singledispatch
+def two_color(graph):
+    """Compute a two-coloring of a directed graph
+
+    If a two coloring is not possible for the input graph (meaning it is not
+    bipartite), ``None`` is returned.
+
+    :param graph: The graph to find the coloring for
+    :returns: If a coloring is possible return a dictionary of node indices to the color as an integer (0 or 1)
+    :rtype: dict
+    """
+
+
+two_color.register(PyDiGraph, digraph_two_color)
+two_color.register(PyGraph, graph_two_color)
+
+
+@functools.singledispatch
+def is_bipartite(graph):
+    """Determine if a given graph is bipartite
+
+    :param graph: The graph to check if it's bipartite
+    :returns: ``True`` if the graph is bipartite and ``False`` if it is not
+    :rtype: bool
+    """
+
+
+is_bipartite.register(PyDiGraph, digraph_is_bipartite)
+is_bipartite.register(PyGraph, graph_is_bipartite)

src/coloring.rs

@@ -10,9 +10,9 @@
 // License for the specific language governing permissions and limitations
 // under the License.
 
-use crate::graph;
+use crate::{digraph, graph};
 
-use rustworkx_core::coloring::{greedy_edge_color, greedy_node_color};
+use rustworkx_core::coloring::{greedy_edge_color, greedy_node_color, two_color};
 
 use pyo3::prelude::*;
 use pyo3::types::PyDict;
@@ -88,3 +88,51 @@ pub fn graph_greedy_edge_color(py: Python, graph: &graph::PyGraph) -> PyResult<P
     }
     Ok(out_dict.into())
 }
+
+/// Compute a two-coloring of a graph
+///
+/// If a two coloring is not possible for the input graph (meaning it is not
+/// bipartite), ``None`` is returned.
+///
+/// :param PyGraph graph: The graph to find the coloring for
+///
+/// :returns: If a coloring is possible return a dictionary of node indices to the color as an
+/// integer (0 or 1)
+/// :rtype: dict
+#[pyfunction]
+pub fn graph_two_color(py: Python, graph: &graph::PyGraph) -> PyResult<Option<PyObject>> {
+    match two_color(&graph.graph) {
+        Some(colors) => {
+            let out_dict = PyDict::new(py);
+            for (node, color) in colors {
+                out_dict.set_item(node.index(), color)?;
+            }
+            Ok(Some(out_dict.into()))
+        }
+        None => Ok(None),
+    }
+}
+
+/// Compute a two-coloring of a directed graph
+///
+/// If a two coloring is not possible for the input graph (meaning it is not
+/// bipartite), ``None`` is returned.
+///
+/// :param PyDiGraph graph: The graph to find the coloring for
+///
+/// :returns: If a coloring is possible return a dictionary of node indices to the color as an
+/// integer (0 or 1)
+/// :rtype: dict
+#[pyfunction]
+pub fn digraph_two_color(py: Python, graph: &digraph::PyDiGraph) -> PyResult<Option<PyObject>> {
+    match two_color(&graph.graph) {
+        Some(colors) => {
+            let out_dict = PyDict::new(py);
+            for (node, color) in colors {
+                out_dict.set_item(node.index(), color)?;
+            }
+            Ok(Some(out_dict.into()))
+        }
+        None => Ok(None),
+    }
+}

src/connectivity/mod.rs

@@ -39,6 +39,7 @@ use crate::iterators::{
 };
 use crate::{EdgeType, StablePyGraph};
 
+use rustworkx_core::coloring::two_color;
 use rustworkx_core::connectivity;
 
 /// Return a list of cycles which form a basis for cycles of a given PyGraph
@@ -1035,3 +1036,23 @@ pub fn digraph_isolates(graph: digraph::PyDiGraph) -> NodeIndices {
             .collect(),
     }
 }
+
+/// Determine if a given graph is bipartite
+///
+/// :param PyGraph graph: The graph to check if it's bipartite
+/// :returns: ``True`` if the graph is bipartite and ``False`` if it is not
+/// :rtype: bool
+#[pyfunction]
+pub fn graph_is_bipartite(graph: graph::PyGraph) -> bool {
+    two_color(&graph.graph).is_some()
+}
+
+/// Determine if a given graph is bipartite
+///
+/// :param PyDiGraph graph: The graph to check if it's bipartite
+/// :returns: ``True`` if the graph is bipartite and ``False`` if it is not
+/// :rtype: bool
+#[pyfunction]
+pub fn digraph_is_bipartite(graph: digraph::PyDiGraph) -> bool {
+    two_color(&graph.graph).is_some()
+}

src/lib.rs

@@ -442,6 +442,10 @@ fn rustworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(digraph_astar_shortest_path))?;
     m.add_wrapped(wrap_pyfunction!(graph_greedy_color))?;
     m.add_wrapped(wrap_pyfunction!(graph_greedy_edge_color))?;
+    m.add_wrapped(wrap_pyfunction!(graph_two_color))?;
+    m.add_wrapped(wrap_pyfunction!(digraph_two_color))?;
+    m.add_wrapped(wrap_pyfunction!(graph_is_bipartite))?;
+    m.add_wrapped(wrap_pyfunction!(digraph_is_bipartite))?;
     m.add_wrapped(wrap_pyfunction!(graph_line_graph))?;
     m.add_wrapped(wrap_pyfunction!(graph_tensor_product))?;
     m.add_wrapped(wrap_pyfunction!(digraph_tensor_product))?;