Better way to loop through nested categories?
<p>I have a sqlalchemy categories table that looks like this:</p>

<pre><code>id | parent_id | name
1 | 0 | license
2 | 1 | Digital Media
3 | 1 | Advertising
4 | 2 | Email Marketing
</code></pre>

<p>My goal is to convert this into a nested dictionary or list that I can loop through in a Flask template. The code below works but looks very messy to me. Is there a way I can improve it?</p>

<pre><code>categories = Category.query.all()
roots = [root for root in categories if root.parent_id == 1]
items = []
for root in roots:
    items.append({'root': root.name, 'subcategories': []})
    for category in categories:
        if category.parent_id == root.id:
            items[-1]['subcategories'].append(category.name) 
</code></pre>

<p>If you just want a more concise way of writing the logic you've put forward, look below:</p>

<pre><code>categories = Category.query.all()
items = [{'root': root.name, 'subcategories':
             [category.name for category in categories if category.parent_id == root.id]
         }
    for root in categories if root.parent_id == 1]
</code></pre>

<p>Note that this runs in O(N<sup>2</sup>) where N is the number of categories. If you want this to preform better, you could map <code>parent_id</code>s to categories. For example:</p>

<pre><code>categories = Category.query.all()
category_by_parent_id = {}
for category in categories:
    if category.parent_id not in category_by_parent_id:
        category_by_parent_id[category.parent_id] = [category]
    else:
        category_by_parent_id[category.parent_id].append(category)
items = [{'root': root.name, 'subcategories': category_by_parent_id[root.id] if root.id in category_by_parent_id else []} for root in category_by_parent_id[1]]
</code></pre>

<p>The above work by making each parent_id map to a list of categories that have that parent id. Then we loop through all root categories (the list at parent_id 1 in the dictionary), and set subcategories equal to the list of categories corresponding to the root's id. If the root's id is not in the dictionary, we instead set subcategories to the empty list. Now the runtime of this operation is O(N) because we're only running through the list twice, and O(2N) = O(N). </p>

<p>This approach can be done even better using a default dictionary from collections. <code>default_dict</code> requires a default value to be specified, and all keys that are not currently in the dictionary are assumed to have that default value. For example <code>a = default_dict(0); print(a[1])</code> would print 0 because that is the default value specified for the list, and key 1 does not currently exist. This can be used to simplify the approach from above:</p>

<pre><code>from collection import default_dict

categories = Category.query.all()
category_by_parent_id = default_dict([])
for category in categories:
    category_by_parent_id[category.parent_id].append(category)
items = [{'root': root.name, 'subcategories': category_by_parent_id[root.id]} for root in category_by_parent_id[1]]
</code></pre>

python