Lesson 5

from networkx import Graph, number_of_nodes, random_tree


def encode(tree):
    if number_of_nodes(tree) <= 2:
        return []

    min_leaf = min([v for v in tree.nodes() if tree.degree[v] == 1])
    parent = list(tree.neighbors(min_leaf))[0]
    tree.remove_node(min_leaf)
    return [parent] + encode(tree)


def decode(code, labels, tree):
    if len(labels) == 2:
        tree.add_edge(labels[0], labels[1])
        return tree

    min_leaf = min([l for l in labels if l not in code])
    tree.add_edge(min_leaf, code[0])
    labels.remove(min_leaf)
    code.remove(code[0])
    return decode(code, labels, tree)


tree = random_tree(n=11, seed=138)
print(tree.edges())
tree_code = encode(tree)
print(tree_code)

decoded_tree = decode(tree_code, list(range(len(tree_code) + 2)), Graph())
print(decoded_tree.edges())

[(0, 3), (1, 6), (2, 3), (2, 10), (3, 7), (3, 4), (4, 6), (4, 8), (5, 7), (8, 9)]
[3, 6, 7, 4, 3, 8, 4, 3, 2]
[(0, 3), (3, 7), (3, 4), (3, 2), (1, 6), (6, 4), (5, 7), (4, 8), (9, 8), (2, 10)]