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