1"""Operations on trees."""
2
3from functools import partial
4from itertools import accumulate, chain
5
6import networkx as nx
7
8__all__ = ["join_trees"]
9
10
11# Argument types don't match dispatching, but allow manual selection of backend
12@nx._dispatchable(graphs=None, returns_graph=True)
13def join_trees(rooted_trees, *, label_attribute=None, first_label=0):
14 """Returns a new rooted tree made by joining `rooted_trees`
15
16 Constructs a new tree by joining each tree in `rooted_trees`.
17 A new root node is added and connected to each of the roots
18 of the input trees. While copying the nodes from the trees,
19 relabeling to integers occurs. If the `label_attribute` is provided,
20 the old node labels will be stored in the new tree under this attribute.
21
22 Parameters
23 ----------
24 rooted_trees : list
25 A list of pairs in which each left element is a NetworkX graph
26 object representing a tree and each right element is the root
27 node of that tree. The nodes of these trees will be relabeled to
28 integers.
29
30 label_attribute : str
31 If provided, the old node labels will be stored in the new tree
32 under this node attribute. If not provided, the original labels
33 of the nodes in the input trees are not stored.
34
35 first_label : int, optional (default=0)
36 Specifies the label for the new root node. If provided, the root node of the joined tree
37 will have this label. If not provided, the root node will default to a label of 0.
38
39 Returns
40 -------
41 NetworkX graph
42 The rooted tree resulting from joining the provided `rooted_trees`. The new tree has a root node
43 labeled as specified by `first_label` (defaulting to 0 if not provided). Subtrees from the input
44 `rooted_trees` are attached to this new root node. Each non-root node, if the `label_attribute`
45 is provided, has an attribute that indicates the original label of the node in the input tree.
46
47 Notes
48 -----
49 Trees are stored in NetworkX as NetworkX Graphs. There is no specific
50 enforcement of the fact that these are trees. Testing for each tree
51 can be done using :func:`networkx.is_tree`.
52
53 Graph, edge, and node attributes are propagated from the given
54 rooted trees to the created tree. If there are any overlapping graph
55 attributes, those from later trees will overwrite those from earlier
56 trees in the tuple of positional arguments.
57
58 Examples
59 --------
60 Join two full balanced binary trees of height *h* to get a full
61 balanced binary tree of depth *h* + 1::
62
63 >>> h = 4
64 >>> left = nx.balanced_tree(2, h)
65 >>> right = nx.balanced_tree(2, h)
66 >>> joined_tree = nx.join_trees([(left, 0), (right, 0)])
67 >>> nx.is_isomorphic(joined_tree, nx.balanced_tree(2, h + 1))
68 True
69
70 """
71 if not rooted_trees:
72 return nx.empty_graph(1)
73
74 # Unzip the zipped list of (tree, root) pairs.
75 trees, roots = zip(*rooted_trees)
76
77 # The join of the trees has the same type as the type of the first tree.
78 R = type(trees[0])()
79
80 lengths = (len(tree) for tree in trees[:-1])
81 first_labels = list(accumulate(lengths, initial=first_label + 1))
82
83 new_roots = []
84 for tree, root, first_node in zip(trees, roots, first_labels):
85 new_root = first_node + list(tree.nodes()).index(root)
86 new_roots.append(new_root)
87
88 # Relabel the nodes so that their union is the integers starting at first_label.
89 relabel = partial(
90 nx.convert_node_labels_to_integers, label_attribute=label_attribute
91 )
92 new_trees = [
93 relabel(tree, first_label=first_label)
94 for tree, first_label in zip(trees, first_labels)
95 ]
96
97 # Add all sets of nodes and edges, attributes
98 for tree in new_trees:
99 R.update(tree)
100
101 # Finally, join the subtrees at the root. We know first_label is unused by
102 # the way we relabeled the subtrees.
103 R.add_node(first_label)
104 R.add_edges_from((first_label, root) for root in new_roots)
105
106 return R