rustworkx_core/generators/
complete_graph.rs

1// Licensed under the Apache License, Version 2.0 (the "License"); you may
2// not use this file except in compliance with the License. You may obtain
3// a copy of the License at
4//
5//     http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
9// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
10// License for the specific language governing permissions and limitations
11// under the License.
12
13use petgraph::data::{Build, Create};
14use petgraph::visit::{Data, GraphProp, NodeIndexable};
15
16use super::utils::get_num_nodes;
17use super::InvalidInputError;
18
19/// Generate a complete graph
20///
21/// Arguments:
22///
23/// * `num_nodes` - The number of nodes to create a complete graph for. Either this or
24///   `weights` must be specified. If both this and `weights` are specified, `weights`
25///   will take priority and this argument will be ignored
26/// * `weights` - A `Vec` of node weight objects.
27/// * `default_node_weight` - A callable that will return the weight to use
28///   for newly created nodes. This is ignored if `weights` is specified.
29/// * `default_edge_weight` - A callable that will return the weight object
30///   to use for newly created edges.
31///
32/// # Example
33/// ```rust
34/// use rustworkx_core::petgraph;
35/// use rustworkx_core::generators::complete_graph;
36/// use rustworkx_core::petgraph::visit::EdgeRef;
37///
38/// let g: petgraph::graph::UnGraph<(), ()> = complete_graph(
39///     Some(4),
40///     None,
41///     || {()},
42///     || {()},
43/// ).unwrap();
44/// assert_eq!(
45///     vec![(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)],
46///     g.edge_references()
47///         .map(|edge| (edge.source().index(), edge.target().index()))
48///         .collect::<Vec<(usize, usize)>>(),
49/// )
50/// ```
51pub fn complete_graph<G, T, F, H, M>(
52    num_nodes: Option<usize>,
53    weights: Option<Vec<T>>,
54    mut default_node_weight: F,
55    mut default_edge_weight: H,
56) -> Result<G, InvalidInputError>
57where
58    G: Build + Create + Data<NodeWeight = T, EdgeWeight = M> + NodeIndexable + GraphProp,
59    F: FnMut() -> T,
60    H: FnMut() -> M,
61{
62    if weights.is_none() && num_nodes.is_none() {
63        return Err(InvalidInputError {});
64    }
65    let node_len = get_num_nodes(&num_nodes, &weights);
66    let mut graph = G::with_capacity(node_len, node_len);
67    if node_len == 0 {
68        return Ok(graph);
69    }
70
71    match weights {
72        Some(weights) => {
73            for weight in weights {
74                graph.add_node(weight);
75            }
76        }
77        None => {
78            for _ in 0..node_len {
79                graph.add_node(default_node_weight());
80            }
81        }
82    };
83    for i in 0..node_len - 1 {
84        for j in i + 1..node_len {
85            let node_i = graph.from_index(i);
86            let node_j = graph.from_index(j);
87            graph.add_edge(node_i, node_j, default_edge_weight());
88            if graph.is_directed() {
89                graph.add_edge(node_j, node_i, default_edge_weight());
90            }
91        }
92    }
93    Ok(graph)
94}
95
96#[cfg(test)]
97mod tests {
98    use crate::generators::complete_graph;
99    use crate::generators::InvalidInputError;
100    use crate::petgraph::graph::{DiGraph, NodeIndex, UnGraph};
101    use crate::petgraph::visit::EdgeRef;
102
103    #[test]
104    fn test_directed_complete_graph() {
105        let g: DiGraph<(), ()> = complete_graph(Some(10), None, || (), || ()).unwrap();
106        assert_eq!(g.node_count(), 10);
107        assert_eq!(g.edge_count(), 90);
108        let mut elist = vec![];
109        for i in 0..10 {
110            for j in i..10 {
111                if i != j {
112                    elist.push((i, j));
113                    elist.push((j, i));
114                }
115            }
116        }
117        assert_eq!(
118            elist,
119            g.edge_references()
120                .map(|edge| (edge.source().index(), edge.target().index()))
121                .collect::<Vec<(usize, usize)>>(),
122        );
123    }
124
125    #[test]
126    fn test_directed_complete_graph_weights() {
127        let g: DiGraph<usize, ()> =
128            complete_graph(None, Some(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), || 4, || ()).unwrap();
129        assert_eq!(g.node_count(), 10);
130        assert_eq!(g.edge_count(), 90);
131        let mut elist = vec![];
132        for i in 0..10 {
133            for j in i..10 {
134                if i != j {
135                    elist.push((i, j));
136                    elist.push((j, i));
137                }
138            }
139            assert_eq!(*g.node_weight(NodeIndex::new(i)).unwrap(), i);
140        }
141        assert_eq!(
142            elist,
143            g.edge_references()
144                .map(|edge| (edge.source().index(), edge.target().index()))
145                .collect::<Vec<(usize, usize)>>(),
146        );
147    }
148
149    #[test]
150    fn test_compete_graph_error() {
151        match complete_graph::<DiGraph<(), ()>, (), _, _, ()>(None, None, || (), || ()) {
152            Ok(_) => panic!("Returned a non-error"),
153            Err(e) => assert_eq!(e, InvalidInputError),
154        };
155    }
156
157    #[test]
158    fn test_complete_graph() {
159        let g: UnGraph<(), ()> = complete_graph(Some(10), None, || (), || ()).unwrap();
160        assert_eq!(g.node_count(), 10);
161        assert_eq!(g.edge_count(), 45);
162        let mut elist = vec![];
163        for i in 0..10 {
164            for j in i..10 {
165                if i != j {
166                    elist.push((i, j));
167                }
168            }
169        }
170        assert_eq!(
171            elist,
172            g.edge_references()
173                .map(|edge| (edge.source().index(), edge.target().index()))
174                .collect::<Vec<(usize, usize)>>(),
175        );
176    }
177
178    #[test]
179    fn test_complete_graph_weights() {
180        let g: UnGraph<usize, ()> =
181            complete_graph(None, Some(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), || 4, || ()).unwrap();
182        assert_eq!(g.node_count(), 10);
183        assert_eq!(g.edge_count(), 45);
184        let mut elist = vec![];
185        for i in 0..10 {
186            for j in i..10 {
187                if i != j {
188                    elist.push((i, j));
189                }
190            }
191            assert_eq!(*g.node_weight(NodeIndex::new(i)).unwrap(), i);
192        }
193        assert_eq!(
194            elist,
195            g.edge_references()
196                .map(|edge| (edge.source().index(), edge.target().index()))
197                .collect::<Vec<(usize, usize)>>(),
198        );
199    }
200}
rustworkx_core/generators/complete_graph.rs

rustworkx_core/generators/
complete_graph.rs
