• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""Generic tools for working with trees."""
2
3from math import ceil, log
4
5
6def build_n_ary_tree(leaves, n):
7    """Build N-ary tree from sequence of leaf nodes.
8
9    Return a list of lists where each non-leaf node is a list containing
10    max n nodes.
11    """
12    if not leaves:
13        return []
14
15    assert n > 1
16
17    depth = ceil(log(len(leaves), n))
18
19    if depth <= 1:
20        return list(leaves)
21
22    # Fully populate complete subtrees of root until we have enough leaves left
23    root = []
24    unassigned = None
25    full_step = n ** (depth - 1)
26    for i in range(0, len(leaves), full_step):
27        subtree = leaves[i : i + full_step]
28        if len(subtree) < full_step:
29            unassigned = subtree
30            break
31        while len(subtree) > n:
32            subtree = [subtree[k : k + n] for k in range(0, len(subtree), n)]
33        root.append(subtree)
34
35    if unassigned:
36        # Recurse to fill the last subtree, which is the only partially populated one
37        subtree = build_n_ary_tree(unassigned, n)
38        if len(subtree) <= n - len(root):
39            # replace last subtree with its children if they can still fit
40            root.extend(subtree)
41        else:
42            root.append(subtree)
43        assert len(root) <= n
44
45    return root
46