chatExcel

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