File size: 1,269 Bytes
bbae060 | 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 46 | """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
|