Add binomial tree graph generator

docs/source/api.rst

@@ -30,6 +30,7 @@ Generators
     retworkx.generators.directed_mesh_graph
     retworkx.generators.grid_graph
     retworkx.generators.directed_grid_graph
+    retworkx.generators.binomial_tree_graph
 
 Random Circuit Functions
 ------------------------

releasenotes/notes/add-binomial-tree-graph-generator-1f6ff6ba3809b901.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Added a new generator for constructing a binomial tree graph (:func:`retworkx.generators.binomial_tree_graph`).

src/generators.rs

@@ -10,11 +10,13 @@
 // License for the specific language governing permissions and limitations
 // under the License.
 
+use std::convert::TryInto;
 use std::iter;
 
 use petgraph::algo;
 use petgraph::graph::NodeIndex;
 use petgraph::stable_graph::{StableDiGraph, StableUnGraph};
+use petgraph::visit::{EdgeRef, IntoEdgeReferences};
 
 use pyo3::exceptions::PyIndexError;
 use pyo3::prelude::*;
@@ -985,6 +987,101 @@ pub fn directed_grid_graph(
     })
 }
 
+/// Generate an undirected binomial tree of order n recursively.
+///
+/// :param int order: Order of the binomial tree.
+/// :param list weights: A list of node weights. It must have 2**order values.
+/// :param bool multigraph: When set to False the output
+///     :class:`~retworkx.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.
+///
+/// :returns: A binomail tree with 2^n vertices and 2^n - 1 edges.
+/// :rtype: PyGraph
+/// :raises IndexError: If neither ``order`` or ``weights`` are specified
+///
+/// .. jupyter-execute::
+///
+///   import os
+///   import tempfile
+///
+///   import pydot
+///   from PIL import Image
+///
+///   import retworkx.generators
+///
+///   graph = retworkx.generators.binomial_tree_graph(4)
+///   dot_str = graph.to_dot(
+///       lambda node: dict(
+///           color='black', fillcolor='lightblue', style='filled'))
+///   dot = pydot.graph_from_dot_data(dot_str)[0]
+///
+///   with tempfile.TemporaryDirectory() as tmpdirname:
+///       tmp_path = os.path.join(tmpdirname, 'dag.png')
+///       dot.write_png(tmp_path)
+///       image = Image.open(tmp_path)
+///       os.remove(tmp_path)
+///   image
+///
+#[pyfunction(multigraph = true)]
+#[text_signature = "(/, order=None, weights=None, multigraph=True)"]
+pub fn binomial_tree_graph(
+    py: Python,
+    order: usize,
+    weights: Option<Vec<PyObject>>,
+    multigraph: bool,
+) -> PyResult<graph::PyGraph> {
+    let mut graph = StableUnGraph::<PyObject, PyObject>::default();
+
+    // try_into is used to convert from usize to u32
+    let num_nodes = usize::pow(2, order.try_into().unwrap());
+
+    let nodes: Vec<NodeIndex> = match weights {
+        Some(weights) => {
+            let mut node_list: Vec<NodeIndex> = Vec::new();
+
+            for weight in weights {
+                let index = graph.add_node(weight);
+                node_list.push(index);
+            }
+
+            node_list
+        }
+
+        None => (0..num_nodes).map(|_| graph.add_node(py.None())).collect(),
+    };
+
+    let mut n = 1;
+
+    for _ in 0..order {
+        let edges: Vec<(NodeIndex, NodeIndex)> = graph
+            .edge_references()
+            .map(|e| (e.source(), e.target()))
+            .collect();
+
+        for (source, target) in edges {
+            let source_index = source.index();
+            let target_index = target.index();
+
+            graph.add_edge(
+                nodes[source_index + n],
+                nodes[target_index + n],
+                py.None(),
+            );
+        }
+
+        graph.add_edge(nodes[0], nodes[n], py.None());
+
+        n *= 2;
+    }
+
+    Ok(graph::PyGraph {
+        graph,
+        node_removed: false,
+        multigraph,
+    })
+}
+
 #[pymodule]
 pub fn generators(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(cycle_graph))?;
@@ -997,5 +1094,6 @@ pub fn generators(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(directed_mesh_graph))?;
     m.add_wrapped(wrap_pyfunction!(grid_graph))?;
     m.add_wrapped(wrap_pyfunction!(directed_grid_graph))?;
+    m.add_wrapped(wrap_pyfunction!(binomial_tree_graph))?;
     Ok(())
 }

tests/generators/test_binomial_tree.py

@@ -0,0 +1,35 @@
+# 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.
+
+import unittest
+
+import retworkx
+
+
+class TestBinomailTreeGraph(unittest.TestCase):
+
+    def test_binomial_tree_graph(self):
+        for n in range(10):
+            graph = retworkx.generators.binomial_tree_graph(n)
+            self.assertEqual(len(graph), 2**n)
+            self.assertEqual(len(graph.edges()), 2**n - 1)
+
+    def test_binomial_tree_graph_weights(self):
+        graph = retworkx.generators.binomial_tree_graph(
+            2, weights=list(range(4)))
+        self.assertEqual(len(graph), 4)
+        self.assertEqual([x for x in range(4)], graph.nodes())
+        self.assertEqual(len(graph.edges()), 3)
+
+    def test_binomial_tree_no_order(self):
+        with self.assertRaises(TypeError):
+            retworkx.generators.binomial_tree_graph(weights=list(range(4)))