Enhancements change default join trees 6947 (networkx#6948)

* change default for new function join_trees networkx#6947 formatting

* regression fix

* post review changes

* Make new kwargs keyword-only.

cf. networkxgh-6956.

* shift label tests to test_basic and inline 1-liner function.

---------

Co-authored-by: Omkar Yadav <[email protected]>
Co-authored-by: Ross Barnowski <[email protected]>
Co-authored-by: Dan Schult <[email protected]>

Author: 4 people

networkx/algorithms/tree/coding.py

@@ -195,7 +195,7 @@ def _make_tree(sequence):
         # For a nonempty sequence, get the subtrees for each child
         # sequence and join all the subtrees at their roots. After
         # joining the subtrees, the root is node 0.
-        return nx.tree.join([(_make_tree(child), 0) for child in sequence])
+        return nx.tree.join_trees([(_make_tree(child), 0) for child in sequence])
 
     # Make the tree and remove the `is_root` node attribute added by the
     # helper function.

networkx/algorithms/tree/operations.py

@@ -31,14 +31,14 @@ def join(rooted_trees, label_attribute=None):
     return join_trees(rooted_trees, label_attribute=label_attribute)
 
 
-def join_trees(rooted_trees, label_attribute=None):
+def join_trees(rooted_trees, *, label_attribute=None, first_label=0):
     """Returns a new rooted tree made by joining `rooted_trees`
 
     Constructs a new tree by joining each tree in `rooted_trees`.
     A new root node is added and connected to each of the roots
     of the input trees. While copying the nodes from the trees,
-    relabeling to integers occurs and the old name stored as an
-    attribute of the new node in the returned graph.
+    relabeling to integers occurs. If the `label_attribute` is provided,
+    the old node labels will be stored in the new tree under this attribute.
 
     Parameters
     ----------
@@ -50,17 +50,20 @@ def join_trees(rooted_trees, label_attribute=None):
 
     label_attribute : str
         If provided, the old node labels will be stored in the new tree
-        under this node attribute. If not provided, the node attribute
-        ``'_old'`` will store the original label of the node in the
-        rooted trees given in the input.
+        under this node attribute. If not provided, the original labels
+        of the nodes in the input trees are not stored.
+
+    first_label : int, optional (default=0)
+        Specifies the label for the new root node. If provided, the root node of the joined tree
+        will have this label. If not provided, the root node will default to a label of 0.
 
     Returns
     -------
     NetworkX graph
-        The rooted tree whose subtrees are the given rooted trees. The
-        new root node is labeled 0. Each non-root node has an attribute,
-        as described under the keyword argument ``label_attribute``,
-        that indicates the label of the original node in the input tree.
+        The rooted tree resulting from joining the provided `rooted_trees`. The new tree has a root node
+        labeled as specified by `first_label` (defaulting to 0 if not provided). Subtrees from the input
+        `rooted_trees` are attached to this new root node. Each non-root node, if the `label_attribute`
+        is provided, has an attribute that indicates the original label of the node in the input tree.
 
     Notes
     -----
@@ -86,42 +89,38 @@ def join_trees(rooted_trees, label_attribute=None):
         True
 
     """
-    if len(rooted_trees) == 0:
+    if not rooted_trees:
         return nx.empty_graph(1)
 
     # Unzip the zipped list of (tree, root) pairs.
     trees, roots = zip(*rooted_trees)
 
-    # The join of the trees has the same type as the type of the first
-    # tree.
+    # The join of the trees has the same type as the type of the first tree.
     R = type(trees[0])()
 
-    # Relabel the nodes so that their union is the integers starting at 1.
-    if label_attribute is None:
-        label_attribute = "_old"
+    lengths = (len(tree) for tree in trees[:-1])
+    first_labels = list(accumulate(lengths, initial=first_label + 1))
+
+    new_roots = []
+    for tree, root, first_node in zip(trees, roots, first_labels):
+        new_root = first_node + list(tree.nodes()).index(root)
+        new_roots.append(new_root)
+
+    # Relabel the nodes so that their union is the integers starting at first_label.
     relabel = partial(
         nx.convert_node_labels_to_integers, label_attribute=label_attribute
     )
-    lengths = (len(tree) for tree in trees[:-1])
-    first_labels = chain([0], accumulate(lengths))
-    trees = [
-        relabel(tree, first_label=first_label + 1)
+    new_trees = [
+        relabel(tree, first_label=first_label)
         for tree, first_label in zip(trees, first_labels)
     ]
 
-    # Get the relabeled roots.
-    roots = [
-        next(v for v, d in tree.nodes(data=True) if d.get(label_attribute) == root)
-        for tree, root in zip(trees, roots)
-    ]
-
     # Add all sets of nodes and edges, attributes
-    for tree in trees:
+    for tree in new_trees:
         R.update(tree)
 
-    # Finally, join the subtrees at the root. We know 0 is unused by the
-    # way we relabeled the subtrees.
-    R.add_node(0)
-    R.add_edges_from((0, root) for root in roots)
+    # Finally, join the subtrees at the root. We know first_label is unused by the way we relabeled the subtrees.
+    R.add_node(first_label)
+    R.add_edges_from((first_label, root) for root in new_roots)
 
     return R

networkx/algorithms/tree/tests/test_operations.py

@@ -4,10 +4,10 @@
 from networkx.utils import edges_equal, nodes_equal
 
 
-def _check_label_attribute(input_trees, res_tree, label_attribute="_old"):
+def _check_custom_label_attribute(input_trees, res_tree, label_attribute):
     res_attr_dict = nx.get_node_attributes(res_tree, label_attribute)
     res_attr_set = set(res_attr_dict.values())
-    input_label = (list(tree[0].nodes()) for tree in input_trees)
+    input_label = (tree for tree, root in input_trees)
     input_label_set = set(chain.from_iterable(input_label))
     return res_attr_set == input_label_set
 
@@ -23,18 +23,31 @@ def test_single():
     """Joining just one tree yields a tree with one more node."""
     T = nx.empty_graph(1)
     trees = [(T, 0)]
-    actual = nx.join_trees(trees)
+    actual_with_label = nx.join_trees(trees, label_attribute="custom_label")
     expected = nx.path_graph(2)
-    assert nodes_equal(list(expected), list(actual))
-    assert edges_equal(list(expected.edges()), list(actual.edges()))
-    assert _check_label_attribute(trees, actual)
+    assert nodes_equal(list(expected), list(actual_with_label))
+    assert edges_equal(list(expected.edges()), list(actual_with_label.edges()))
 
 
 def test_basic():
     """Joining multiple subtrees at a root node."""
     trees = [(nx.full_rary_tree(2, 2**2 - 1), 0) for i in range(2)]
-    label_attribute = "old_values"
-    actual = nx.join_trees(trees, label_attribute)
     expected = nx.full_rary_tree(2, 2**3 - 1)
+    actual = nx.join_trees(trees, label_attribute="old_labels")
     assert nx.is_isomorphic(actual, expected)
-    assert _check_label_attribute(trees, actual, label_attribute)
+    assert _check_custom_label_attribute(trees, actual, "old_labels")
+
+    actual_without_label = nx.join_trees(trees)
+    assert nx.is_isomorphic(actual_without_label, expected)
+    # check that no labels were stored
+    assert all(not data for _, data in actual_without_label.nodes(data=True))
+
+
+def test_first_label():
+    """Test the functionality of the first_label argument."""
+    T1 = nx.path_graph(3)
+    T2 = nx.path_graph(2)
+    actual = nx.join_trees([(T1, 0), (T2, 0)], first_label=10)
+    expected_nodes = set(range(10, 16))
+    assert set(actual.nodes()) == expected_nodes
+    assert set(actual.neighbors(10)) == {11, 14}