Making a tree out of a 2 column list

[bearophileHUGS at lycos.com]

The solution I have found is quite similar to the martinskou at gmail.com
one (it uses a defaultdict, and it yields the result lazily), but it
lacks the quite useful max_depth (it can be added):

from collections import defaultdict

def finder(el, stor):
    if el in stor:
        for v in stor[el]:
            yield v
            for v2 in finder(v, stor):
                yield v2

data = [[131, 2], [335, 6], [6, 7], [9, 8], [131, 10], [99, 131], [10,
5]]

storage = defaultdict(list)
for k, v in data:
    storage[k].append(v)

test_data = ([131, [2, 10, 5]],
             [335, [6, 7]],
             [9, [8]],
            )

print storage
for k, vals in test_data:
    assert set(finder(k, storage)) == set(vals)

Bye,
bearophile