| """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 | |