Printing the hierarchical representation, and finding all ascendants for a node, in a self-referential table in Python with recursion

Question

I have a self referential database table:

CREATE TABLE category (
    id INT PRIMARY KEY, 
    name VARCHAR,
    parent_id INT REFERENCES category (id)
);

And the following INSERTS:

INSERT INTO category (id, name, parent_id) VALUES (1, 'Programming', null);
INSERT INTO category (id, name, parent_id) VALUES (2, 'Python', 1);
INSERT INTO category (id, name, parent_id) VALUES (3, 'Metaclasses', 2);
INSERT INTO category (id, name, parent_id) VALUES (4, 'Java', 1);
INSERT INTO category (id, name, parent_id) VALUES (5, 'Reflection', 4);
INSERT INTO category (id, name, parent_id) VALUES (6, 'Homebrewing', null);
INSERT INTO category (id, name, parent_id) VALUES (7, 'Yeast', 6);

My goal is to output the table in the following hierarchical format:

Programming
    Python
        Metaclasses
    Java
        Reflection
Homebrewing
    Yeast

Here is my solution:

class Category(object):
    #This is actually a SQLAlchemy table
    #I've included this for others to reproduce easier
    def __init__(self, id, name, parent_id=None):
        self.id = id
        self.name = name
        self.parent_id = parent_id
    def __repr__(self):
        return "Category(id={}, name='{}', parent_id={})".format(self.id, self.name, self.parent_id)

def _append_tree(category, tree, parent=None):
    # The only reason to include parent is for the find_ascendant generator
    tree.append({'parent': parent, 'category': category, 'children': []})

# This looks fine to me. Can it be better?
def add_node(category, tree):
    if category.parent_id is None:
        _append_tree(category, tree)
        return 1
    for child in tree:
        if category.parent_id == child['category'].id:
            # child obj has a ref to his parent
            _append_tree(category, child['children'], child)
            return 1
        if add_node(category, child['children']):
            return 1
    return 0

def print_tree(tree, depth=0):
    for child in tree:
        print '    ' * depth + child['category'].name
        print_tree(child['children'], depth + 1)

# This works, but looks bad; any suggestions?
def iter_tree(tree, depth=0):
    for child in tree:
        yield depth, child['category']
        for i in iter_tree(child['children'], depth + 1):
            yield i

This is the function for finding all ascendants for a leaf.

def _find_ascendants(child):
    yield child['category']
    if child['parent'] is not None:
        for i in _find_ascendants(child['parent']):
            yield i

def find_ascendants(category, tree):
    if category.parent_id is None:
        return
    for child in tree:
        if category.parent_id == child['category'].id:
            # We've found the parent category. Return the parent's parent
            for i in _find_ascendants(child):
                yield i
                # No need to check any more nodes
                return
        # No match; try children
        for i in find_ascendants(category, child['children']):
            yield i

This has the entire flow to show that it works:

tree = []

# The order by returns NULLS FIRST, so root nodes will always be populated before children
for category in session.query(Category).order_by(Category.parent_id):
    add_node(category, tree)

print_tree(tree)

for depth, child in iter_tree(tree):
    print '----' * depth, child.name

for i in find_ascendants(Category(id=3, name='Metaclasses', parent_id=2), tree):
    print i

# or, to reproduce without a database:
tree = []

categories = [
    Category(1, 'Programming', None),
    Category(6, 'Homebrewing', None),
    Category(2, 'Python', 1),
    Category(4, 'Java', 1),
    Category(3, 'Metaclasses', 2),
    Category(5, 'Reflection', 5),
    Category(7, 'Yeast', 6)

]

for category in categories:
     add_node(category, tree)

print_tree(tree)

for depth, child in iter_tree(tree):
    print '----' * depth, child.name

for i in find_ascendants(Category(id=3, name='Metaclasses', parent_id=2), tree):
    print i

I am notably concerned about the iter_tree generator and the find_ascendants generator. I would also appreciate any suggestions to improve add_node.

Here is an example of how I want to use the find_ascendants generator.

The user creates a post and applies two categories to it, Maybe Metaclasses and Yeast, for example. I want to find all the parent categories of Metaclasses and Yeast and then apply them to the post as well.

# User provides this
category_ids = [3, 7]  # Metaclasses and Yeast

# Construct tree
tree = []
for category in db.session.query(Category).order_by(Category.parent_id):
    add_node(category, tree)

all_category_ids = set(category_ids)

for category in db.session.query(Category).filter(Category.id.in_(category_ids)):
    ascendants = find_ascendants(category, tree)
    for ascendant in ascendants:
        all_category_ids.add(ascendant.id)

assert all_category_ids == set([1,2,3,6,7])
# Do something with all_category_ids here

Answer 1

The overall code is relatively good, but there are 2 things that may help improve add_node and find_ascendants:

• Using dictionnaries instead of lists;
• Using an helper function to find a node with the given ID.

If you really want to retain the order of inserted elements, you could still use an OrderedDict from collections.

def _get_node_with_id(id, tree):
    try:
        return tree[id]
    except KeyError:
        for child_tree in tree.values():
            node = _get_node_with_id(id, child_tree['children'])
            if node is not None:
                return node

def add_node(category, tree):
    parent = category.parent_id
    if parent is None:
        node = tree
    else:
        node = _get_node_with_id(parent, tree)['children']
    node[category.id] = {'category': category, 'children': {}}

def find_ascendants(category, tree):
    id = category.id
    while id not in tree:
        node = _get_node_with_id(id)
        id = node['category'].parent_id
        yield id

You’ll note that I indexed the categories with their ID, I feel it simplifies a lot the search.

The iter_tree (which I find pretty clean as it stand) should be adapted a bit:

def iter_tree(tree, depth=0):
    for node in tree.values():
        yield depth, node['category']
        for i in iter_tree(node['children'], depth + 1):
            yield i

And you should have print_tree use it instead of somewhat duplicating the logic:

def print_tree(tree):
    for depth, node in iter_tree(tree):
        print '    ' * depth + node['category'].name

And, last thing, your Category class is pretty much:

Category = collections.namedtuple('Category', 'id name parent_id')

---

Using _get_node_with_id like that gives a neat writting but you lose informations about the recursion when you have a simple return. You can turn it into a generator to reduce the time complexity of find_ascendants but its writting will be more complicated:

def _get_node_with_id(id, tree):
    try:
        yield tree[id]
    except KeyError:
        for child_tree in tree.values():
            nodes = _get_node_with_id(id, child_tree['children'])
            try:
                # Check if the path is right
                yield next(nodes)
            except StopIteration:
                pass
            else:
                for node in nodes:
                    yield node  # yield previous items on the path
                yield child_tree  # yield current item
                return  # stop checking as we found the right path

def add_node(category, tree):
    parent = category.parent_id
    if parent is None:
        node = tree
    else:
        nodes = _get_node_with_id(parent, tree)
        node = next(nodes)['children']  # Take only the leaf node
    node[category.id] = {'category': category, 'children': {}}

def find_ascendants(category, tree):
    nodes = _get_node_with_id(id)
    next(nodes)  # Skip ourselves as we only want ascendants
    for node in nodes:
        yield node['category'].id

The idea, with this version of _get_node_with_id is based on the same principle than the previous one: we perform depth-first search on the tree until we reach the node with the right id. At this point we yield it (the dictionary containing category + children), but because of it, we cannot check for None to know that the path was wrong and move on to the next one.

So we try to iterate over the result instead: we can, it was the right path, otherwise we need to try a new one. If it was the right path, yield every nodes already encountered on this path, add the parent node and return to stop searching the tree.

The "ugly" thing is that we need to manually iterate over the first item to know for sure there is something to iterate over, since a for loop would happily do nothing if there was nothing to iterate over. Manually calling next let us catch the StopIteration exception and know that, since there was no node found, we need to check the next path.