# -*- coding: utf-8 -*-
# python

# Programing Challenge: Constructing a Tree Given Its Edges
# http://xahlee.info/perl-python/python_construct_tree_from_edge.html

import random
import json                     # pretty print nested hash
import collections

def randomTreeEdges(max_nodes):
    id_pool = list(range(max_nodes))
    random.shuffle(id_pool)
    root = id_pool.pop()
    first_child = id_pool.pop()
    edges = [[first_child, root]]
    while len(id_pool) != 0:
        if random.random() > .1: # add child
            edges.append([id_pool.pop(), random.choice(edges)[0]])
        else:                   # add sibling
            edges.append([id_pool.pop(), random.choice(edges)[1]])
    random.shuffle(edges)
    return edges

def construct_trees_by_TingYu(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}

def construct_tree_by_Ian(edges):
    forest = {}
    # Populate the forest with depth-1 trees, with placeholders for subtrees.
    for (child, parent) in edges:
        if parent not in forest:
            forest[parent] = {}
        forest[parent][child] = {}
    if not forest:
        # Graph is empty; an edge list cannot represent a single-node tree.
        return {}
    roots = set(forest.keys())
    # Replace the placeholders with corresponding trees from the forest.
    for tree in forest.values():
        roots -= tree.keys()
        for child in tree:
            if child in forest:
                tree[child] = forest[child]
    # Make sure we have a single root.  No check is made for cycles.
    if len(roots) == 0 or len(roots) > 1:
        raise ValueError("Graph is not a tree")
    # The tree located at the actual root contains the entire tree.
    root = roots.pop()
    tree = {root: forest[root]}
    return tree

if __name__ == '__main__':
    edges = randomTreeEdges(20000)
    import timeit

    # t2 = construct_trees_by_TingYu(edges)
    # t3 = construct_tree_by_Ian(edges)
    # if t2 == t3:
    #     print( "all's good")

    # print( "edges {}".format(edges))
    # print(json.dumps(t2, indent=1))

    print (timeit.timeit(stmt="construct_trees_by_TingYu(edges)",setup="from __main__ import construct_trees_by_TingYu, edges", number=1))
    print (timeit.timeit(stmt="construct_tree_by_Ian(edges)",setup="from __main__ import construct_tree_by_Ian, edges", number=1))