Skip to content

Commit

Permalink
Add two_color and is_bipartite
Browse files Browse the repository at this point in the history
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 Qiskit#998 and will need to be rebased after
that merges.
  • Loading branch information
mtreinish committed Oct 12, 2023
1 parent 274b004 commit 8cde714
Show file tree
Hide file tree
Showing 8 changed files with 396 additions and 7 deletions.
22 changes: 22 additions & 0 deletions releasenotes/notes/add-bipartite-9df3898a156e799c.yaml
Original file line number Diff line number Diff line change
@@ -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.
159 changes: 154 additions & 5 deletions rustworkx-core/src/coloring.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,90 @@
use std::cmp::Reverse;
use std::hash::Hash;

use crate::connectivity::isolates;
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);
}
}
}
}
colors.extend(isolates(&graph).into_iter().map(|x| (x, 0)));
Some(colors)
}

/// Color a graph using a greedy graph coloring algorithm.
///
/// This function uses a `largest-first` strategy as described in:
Expand Down Expand Up @@ -150,11 +227,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() {
Expand Down Expand Up @@ -201,6 +279,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), 0);
expected_colors.insert(NodeIndex::new(6), 0);
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), 0);
expected_colors.insert(NodeIndex::new(6), 0);
assert_eq!(coloring, expected_colors)
}
}

#[cfg(test)]
Expand Down
31 changes: 31 additions & 0 deletions rustworkx/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
52 changes: 50 additions & 2 deletions src/coloring.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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),
}
}
21 changes: 21 additions & 0 deletions src/connectivity/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
}
4 changes: 4 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))?;
Expand Down
Loading

0 comments on commit 8cde714

Please sign in to comment.