Build a tree from a list of pairs of (parent, child) identifiers.

build_tree.py

from collections import defaultdict

ROOT = -1

def create_tree(pairs):
    """ Given an iterable of (parent, child) identifiers, build a tree
        where each node is a dict whose key is its identifier, and whose
        value is a list of children (which may be empty for leaf nodes).
        Multiple root nodes are supported; all roots are placed under a
        synthetic ROOT node.
    """
    def node(): return [ROOT,[]]
    table = defaultdict(node)
    # Build 2-way mapping between nodes
    for parent, child in pairs:
        table[parent][1].append(child)  # parent - > children
        table[child][0] = parent  # child -> parent

    def follow(parent, childids):
        for c in childids:
            empty = []
            child = {c: empty}
            parent.append(child)
            if c in table:
                follow(empty, table[c][1])
                
    # Recursively fill in the tree
    tree = {ROOT:[]}
    roots = [k for k,v in table.items() if v[0] == ROOT]
    follow(tree[ROOT], roots)   
    return tree

#
# TESTS:
#

def _build_test_tree(nodes):
    def _build(nodes, tree):
        if len(nodes) == 0: return []
        else: return [{nodes[0]: _build(nodes[1:], tree)}]
    return _build(nodes, {})[0]

def assertEquals(actual, expected):
    assert actual == expected, "%s != %s" % (actual, expected)

def test_empty():
    assert create_tree([]) == {ROOT:[]}

def test_one():
    expected = _build_test_tree([ROOT,1,2])
    actual = create_tree([(1,2)])
    assertEquals(actual, expected)

def test_chain():
    expected = _build_test_tree([ROOT,1,2,3])
    actual = create_tree([(1,2),(2,3)])
    assertEquals(actual, expected)

def test_forked_root():
    two = {2: []}
    four = {4: []}
    actual = create_tree([(1,2),(3,4)])
    expected = {ROOT: [{1: [two]},{3: [four]}]}
    assertEquals(actual, expected)

def test_replace_root():
    expected = _build_test_tree([ROOT,1,2,3])
    actual = create_tree([(2,3),(1,2)])
    assertEquals(actual, expected)

def test_splice():
    expected = _build_test_tree([ROOT,1,2,3,4])
    actual = create_tree([(3,4),(2,3),(1,2)])
    assertEquals(actual, expected)

## Fails because the order of the lists may vary, which is OK but hard to test.
## TODO restore this test after changing to {id: 1, name: 'foo', children: []} format...
##def test_order_independent():
##    from random import shuffle
##    pairs = [(1,2),(2,3),(2,4),(1,5),(6,7),(6,8),(8,9)]
##    results = []
##    for i in range(10):
##        shuffle(pairs)
##        results.append(create_tree(pairs))
##    for i in range(9):
##        assert results[i] == results[i+1], "Different input order gives different results: %s != %s" % (results[i], results[i+1])
    
if __name__ == "__main__":
    tests = [(k, v) for k, v in locals().items() if k.startswith("test")]
    print("%s TESTS FOUND" % len(tests))
    for k, v in tests:
        print(k)
        v()
    print("TESTS COMPLETE")