diff --git a/rustworkx-core/src/generators/lollipop_graph.rs b/rustworkx-core/src/generators/lollipop_graph.rs new file mode 100644 index 000000000..afbdfc69e --- /dev/null +++ b/rustworkx-core/src/generators/lollipop_graph.rs @@ -0,0 +1,200 @@ +// 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, GraphProp, NodeIndexable}; + +use super::utils::pairwise; +use super::InvalidInputError; + +/// Generate a lollipop graph +/// +/// Arguments: +/// +/// Generate an undirected lollipop graph where a mesh graph is connected to a +/// path. +/// +/// If neither `num_path_nodes` nor `path_weights` (both described +/// below) are specified then this is equivalent to +/// :func:`~rustworkx.generators.mesh_graph` +/// +/// * `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` is 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::lollipop_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), +/// ]; +/// let g: petgraph::graph::UnGraph<(), ()> = lollipop_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 lollipop_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 + GraphProp, + F: FnMut() -> T, + H: FnMut() -> M, + G::NodeId: Eq + Hash, +{ + if num_mesh_nodes.is_none() && mesh_weights.is_none() { + return Err(InvalidInputError {}); + } + let num_nodes: usize; + if let Some(mesh_nodes) = num_mesh_nodes { + num_nodes = mesh_nodes; + } else { + num_nodes = mesh_weights.as_ref().unwrap().len(); + } + let num_edges = (num_nodes * (num_nodes - 1)) / 2; + let mut graph = G::with_capacity(num_nodes, num_edges); + + 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()); + } + } + 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, + }; + } + } + Ok(graph) +} + +#[cfg(test)] +mod tests { + use crate::generators::lollipop_graph; + use crate::generators::InvalidInputError; + use crate::petgraph::visit::EdgeRef; + #[test] + fn test_lollipop_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), + ]; + let g: petgraph::graph::UnGraph<(), ()> = + lollipop_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_lollipop_none_mesh() { + match lollipop_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 ff1b8e5c6..9ecc58f19 100644 --- a/rustworkx-core/src/generators/mod.rs +++ b/rustworkx-core/src/generators/mod.rs @@ -19,6 +19,7 @@ mod grid_graph; mod heavy_hex_graph; mod heavy_square_graph; mod hexagonal_lattice_graph; +mod lollipop_graph; mod path_graph; mod petersen_graph; mod star_graph; @@ -47,6 +48,7 @@ pub use grid_graph::grid_graph; pub use heavy_hex_graph::heavy_hex_graph; pub use heavy_square_graph::heavy_square_graph; pub use hexagonal_lattice_graph::hexagonal_lattice_graph; +pub use lollipop_graph::lollipop_graph; pub use path_graph::path_graph; pub use petersen_graph::petersen_graph; pub use star_graph::star_graph; diff --git a/src/generators.rs b/src/generators.rs index 9f28288ac..f181d59f5 100644 --- a/src/generators.rs +++ b/src/generators.rs @@ -1332,37 +1332,28 @@ pub fn lollipop_graph( path_weights: Option>, multigraph: bool, ) -> PyResult { - let mut graph = mesh_graph(py, num_mesh_nodes, mesh_weights, multigraph)?; - if num_path_nodes.is_none() && path_weights.is_none() { - return Ok(graph); - } - let meshlen = graph.num_nodes(); - - let path_nodes: Vec = match path_weights { - Some(path_weights) => path_weights - .into_iter() - .map(|node| graph.graph.add_node(node)) - .collect(), - None => (0..num_path_nodes.unwrap()) - .map(|_| graph.graph.add_node(py.None())) - .collect(), - }; - - let pathlen = path_nodes.len(); - if pathlen > 0 { - graph.graph.add_edge( - NodeIndex::new(meshlen - 1), - NodeIndex::new(meshlen), - py.None(), - ); - for (node_a, node_b) in pairwise(path_nodes) { - match node_a { - Some(node_a) => graph.graph.add_edge(node_a, node_b, py.None()), - None => continue, - }; + let default_fn = || py.None(); + let graph: StablePyGraph = match core_generators::lollipop_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", + )) } - } - Ok(graph) + }; + Ok(graph::PyGraph { + graph, + node_removed: false, + multigraph, + attrs: py.None(), + }) } /// Generate a generalized Petersen graph :math:`G(n, k)` with :math:`2n`