construct_trees.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
This code resolves the problem described in
http://xahlee.info/perl-python/python_construct_tree_from_edge.html

"""

import collections
import json


def construct_trees(edges):
    """Given a list of edges [child, parent], return trees. """
    trees = collections.defaultdict(dict)

    for child, parent in edges:
        trees[parent][child] = trees[child]

    # Find roots
    children, parents = zip(*edges)
    roots = set(parents).difference(children)

    return {root: trees[root] for root in roots}


if __name__ == '__main__':
    edges = [[0, 2], [3, 0], [1, 4], [2, 4], [5, 6], [6, 7], [8, 6]]
    print(json.dumps(construct_trees(edges), indent=2))