diff --git a/rustworkx-core/src/generators/barbell_graph.rs b/rustworkx-core/src/generators/barbell_graph.rs new file mode 100644 index 0000000000..e7383f5afc --- /dev/null +++ b/rustworkx-core/src/generators/barbell_graph.rs @@ -0,0 +1,239 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); you may +// not use this file except in compliance with the License. You may obtain +// a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +// License for the specific language governing permissions and limitations +// under the License. + +use std::hash::Hash; + +use petgraph::data::{Build, Create}; +use petgraph::visit::{Data, NodeIndexable, NodeRef}; + +use super::utils::get_num_nodes; +use super::utils::pairwise; +use super::InvalidInputError; + +/// Generate an undirected barbell graph where two identical mesh graphs are +/// connected by a path. +/// +/// .. note:: +/// +/// If neither `num_path_nodes` nor `path_weights` (described below) is not specified +/// then this is equivalent to two mesh graphs joined together. +/// +/// Arguments: +/// +/// * `num_mesh_nodes` - The number of nodes to generate the mesh graph +/// with. Node weights will be None if this is specified. If both +/// `num_mesh_nodes` and `mesh_weights` are set this will be ignored and +/// `mesh_weights` will be used. +/// * `num_path_nodes` - The number of nodes to generate the path +/// with. Node weights will be None if this is specified. If both +/// `num_path_nodes` and `path_weights` are set this will be ignored and +/// `path_weights` will be used. +/// * `mesh_weights` - A list of node weights for the mesh graph. If both +/// `num_mesh_nodes` and `mesh_weights` are set `num_mesh_nodes` will +/// be ignored and `mesh_weights` will be used. +/// * `path_weights` - A list of node weights for the path. If both +/// `num_path_nodes` and `path_weights` are set `num_path_nodes` will +/// be ignored and `path_weights` will be used. +/// * `default_node_weight` - A callable that will return the weight to use +/// for newly created nodes. This is ignored if weights are specified, +/// as the weights from that argument will be used instead. +/// * `default_edge_weight` - A callable that will return the weight object +/// to use for newly created edges. +/// +/// # Example +/// ```rust +/// use rustworkx_core::petgraph; +/// use rustworkx_core::generators::barbell_graph; +/// use rustworkx_core::petgraph::visit::EdgeRef; +/// +/// let expected_edge_list = vec![ +/// (0, 1), +/// (0, 2), +/// (0, 3), +/// (1, 2), +/// (1, 3), +/// (2, 3), +/// (3, 4), +/// (4, 5), +/// (5, 6), +/// (6, 7), +/// (7, 8), +/// (7, 9), +/// (7, 10), +/// (8, 9), +/// (8, 10), +/// (9, 10), +/// ]; +/// let g: petgraph::graph::UnGraph<(), ()> = barbell_graph( +/// Some(4), +/// Some(3), +/// None, +/// None, +/// || {()}, +/// || {()}, +/// ).unwrap(); +/// assert_eq!( +/// expected_edge_list, +/// g.edge_references() +/// .map(|edge| (edge.source().index(), edge.target().index())) +/// .collect::>(), +/// ) +/// ``` +pub fn barbell_graph( + num_mesh_nodes: Option, + num_path_nodes: Option, + mesh_weights: Option>, + path_weights: Option>, + mut default_node_weight: F, + mut default_edge_weight: H, +) -> Result +where + G: Build + Create + Data + NodeIndexable, + F: FnMut() -> T, + H: FnMut() -> M, + G::NodeId: Eq + Hash + NodeRef, + T: Clone, +{ + if num_mesh_nodes.is_none() && mesh_weights.is_none() { + return Err(InvalidInputError {}); + } + let num_nodes = get_num_nodes(&num_mesh_nodes, &mesh_weights); + let num_edges = (num_nodes * (num_nodes - 1)) / 2; + let mesh_weights_2 = mesh_weights.clone(); + + let mut graph = G::with_capacity(num_nodes, num_edges); + + // Create left mesh graph + let mesh_nodes: Vec = match mesh_weights { + Some(mesh_weights) => mesh_weights + .into_iter() + .map(|weight| graph.add_node(weight)) + .collect(), + None => (0..num_mesh_nodes.unwrap()) + .map(|_| graph.add_node(default_node_weight())) + .collect(), + }; + + let meshlen = mesh_nodes.len(); + for i in 0..meshlen - 1 { + for j in i + 1..meshlen { + graph.add_edge(mesh_nodes[i], mesh_nodes[j], default_edge_weight()); + } + } + + // Add path nodes and edges + let path_nodes: Vec = match path_weights { + Some(path_weights) => path_weights + .into_iter() + .map(|weight| graph.add_node(weight)) + .collect(), + None => { + if let Some(num_path) = num_path_nodes { + (0..num_path) + .map(|_| graph.add_node(default_node_weight())) + .collect() + } else { + vec![] + } + } + }; + let pathlen = path_nodes.len(); + if pathlen > 0 { + graph.add_edge( + graph.from_index(meshlen - 1), + graph.from_index(meshlen), + default_edge_weight(), + ); + for (node_a, node_b) in pairwise(path_nodes) { + match node_a { + Some(node_a) => graph.add_edge(node_a, node_b, default_edge_weight()), + None => continue, + }; + } + } + + // Add right mesh graph + let mesh_nodes: Vec = match mesh_weights_2 { + Some(mesh_weights_2) => mesh_weights_2 + .into_iter() + .map(|weight| graph.add_node(weight)) + .collect(), + None => (0..num_mesh_nodes.unwrap()) + .map(|_| graph.add_node(default_node_weight())) + .collect(), + }; + let meshlen = mesh_nodes.len(); + + graph.add_edge( + graph.from_index(meshlen + pathlen - 1), + graph.from_index(meshlen + pathlen), + default_edge_weight(), + ); + for i in 0..meshlen - 1 { + for j in i + 1..meshlen { + graph.add_edge(mesh_nodes[i], mesh_nodes[j], default_edge_weight()); + } + } + Ok(graph) +} + +#[cfg(test)] +mod tests { + use crate::generators::barbell_graph; + use crate::generators::InvalidInputError; + use crate::petgraph::visit::EdgeRef; + + #[test] + fn test_barbell_mesh_path() { + let expected_edge_list = vec![ + (0, 1), + (0, 2), + (0, 3), + (1, 2), + (1, 3), + (2, 3), + (3, 4), + (4, 5), + (5, 6), + (6, 7), + (7, 8), + (7, 9), + (7, 10), + (8, 9), + (8, 10), + (9, 10), + ]; + let g: petgraph::graph::UnGraph<(), ()> = + barbell_graph(Some(4), Some(3), None, None, || (), || ()).unwrap(); + assert_eq!( + expected_edge_list, + g.edge_references() + .map(|edge| (edge.source().index(), edge.target().index())) + .collect::>(), + ); + } + + #[test] + fn test_barbell_none_mesh() { + match barbell_graph::, (), _, _, ()>( + None, + None, + None, + None, + || (), + || (), + ) { + Ok(_) => panic!("Returned a non-error"), + Err(e) => assert_eq!(e, InvalidInputError), + }; + } +} diff --git a/rustworkx-core/src/generators/mod.rs b/rustworkx-core/src/generators/mod.rs index 0b2bd138c7..ff1b8e5c61 100644 --- a/rustworkx-core/src/generators/mod.rs +++ b/rustworkx-core/src/generators/mod.rs @@ -12,6 +12,7 @@ //! This module contains generator functions for building graphs +mod barbell_graph; mod binomial_tree_graph; mod cycle_graph; mod grid_graph; @@ -39,6 +40,7 @@ impl fmt::Display for InvalidInputError { } } +pub use barbell_graph::barbell_graph; pub use binomial_tree_graph::binomial_tree_graph; pub use cycle_graph::cycle_graph; pub use grid_graph::grid_graph; diff --git a/rustworkx-core/src/generators/utils.rs b/rustworkx-core/src/generators/utils.rs index b63ab70438..3eb178fb5f 100644 --- a/rustworkx-core/src/generators/utils.rs +++ b/rustworkx-core/src/generators/utils.rs @@ -10,6 +10,8 @@ // License for the specific language governing permissions and limitations // under the License. +use std::iter; + #[inline] pub fn get_num_nodes(num_nodes: &Option, weights: &Option>) -> usize { if weights.is_some() { @@ -18,3 +20,11 @@ pub fn get_num_nodes(num_nodes: &Option, weights: &Option>) -> num_nodes.unwrap() } } + +pub fn pairwise(right: I) -> impl Iterator, I::Item)> +where + I: IntoIterator + Clone, +{ + let left = iter::once(None).chain(right.clone().into_iter().map(Some)); + left.zip(right) +} diff --git a/src/generators.rs b/src/generators.rs index 4fe607d244..9f28288ace 100644 --- a/src/generators.rs +++ b/src/generators.rs @@ -16,7 +16,6 @@ use std::iter; use petgraph::algo; use petgraph::graph::NodeIndex; use petgraph::prelude::*; -use petgraph::visit::{EdgeRef, IntoEdgeReferences}; use petgraph::Undirected; use pyo3::exceptions::{PyIndexError, PyOverflowError}; @@ -1456,6 +1455,12 @@ pub fn generalized_petersen_graph( /// :class:`~rustworkx.PyGraph` object will not be not be a multigraph and /// won't allow parallel edges to be added. Instead /// calls which would create a parallel edge will update the existing edge. +/// :param list mesh_weights: A list of node weights for the mesh graph. If both +/// ``num_mesh_nodes`` and ``mesh_weights`` are set ``num_mesh_nodes`` will +/// be ignored and ``mesh_weights`` will be used. +/// :param list path_weights: A list of node weights for the path. If both +/// ``num_path_nodes`` and ``path_weights`` are set ``num_path_nodes`` will +/// be ignored and ``path_weights`` will be used. /// /// :returns: The generated barbell graph /// :rtype: PyGraph @@ -1476,62 +1481,27 @@ pub fn barbell_graph( num_mesh_nodes: Option, num_path_nodes: Option, multigraph: bool, + mesh_weights: Option>, + path_weights: Option>, ) -> PyResult { - if num_mesh_nodes.is_none() { - return Err(PyIndexError::new_err("num_mesh_nodes not specified")); - } - - let mut left_mesh = StableUnGraph::::default(); - let mesh_nodes: Vec = (0..num_mesh_nodes.unwrap()) - .map(|_| left_mesh.add_node(py.None())) - .collect(); - let mut nodelen = mesh_nodes.len(); - for i in 0..nodelen - 1 { - for j in i + 1..nodelen { - left_mesh.add_edge(mesh_nodes[i], mesh_nodes[j], py.None()); - } - } - - let right_mesh = left_mesh.clone(); - - if let Some(num_nodes) = num_path_nodes { - let path_nodes: Vec = (0..num_nodes) - .map(|_| left_mesh.add_node(py.None())) - .collect(); - left_mesh.add_edge( - NodeIndex::new(nodelen - 1), - NodeIndex::new(nodelen), - py.None(), - ); - - nodelen += path_nodes.len(); - - for (node_a, node_b) in pairwise(path_nodes) { - match node_a { - Some(node_a) => left_mesh.add_edge(node_a, node_b, py.None()), - None => continue, - }; + let default_fn = || py.None(); + let graph: StablePyGraph = match core_generators::barbell_graph( + num_mesh_nodes, + num_path_nodes, + mesh_weights, + path_weights, + default_fn, + default_fn, + ) { + Ok(graph) => graph, + Err(_) => { + return Err(PyIndexError::new_err( + "num_nodes and weights list not specified", + )) } - } - - for node in right_mesh.node_indices() { - let new_node = &right_mesh[node]; - left_mesh.add_node(new_node.clone_ref(py)); - } - left_mesh.add_edge( - NodeIndex::new(nodelen - 1), - NodeIndex::new(nodelen), - py.None(), - ); - for edge in right_mesh.edge_references() { - let new_source = NodeIndex::new(nodelen + edge.source().index()); - let new_target = NodeIndex::new(nodelen + edge.target().index()); - let weight = edge.weight(); - left_mesh.add_edge(new_source, new_target, weight.clone_ref(py)); - } - + }; Ok(graph::PyGraph { - graph: left_mesh, + graph, node_removed: false, multigraph, attrs: py.None(),