Implement depth first tree traversal with stack in Python

Question

If given a JSON list of locations, for example:

locations = [
  {"id": 1, "name": "San Francisco Bay Area", "parent_id": None},
  {“id": 2, "name": "San Jose", "parent_id": 3},
  {"id": 3, "name": "South Bay", "parent_id": 1},
  {"id": 4, "name": "San Francisco", "parent_id": 1},
  {"id": 5, "name": "Manhattan", "parent_id": 6},
  {"id": 6, "name": "New York", "parent_id": None}
]

I want to be able to generate a list of the locations, with sub-locations grouped under their parents, and in alphabetical order, also indenting sub-locations with a hyphen. Each level of depth should be alphabetically sorted, and there can be up to 5 levels of depth. So the output for the above would be:

New York
-Manhattan
San Francisco Bay Area
-San Francisco
-South Bay
--San Jose

It seems like it would make sense to go through the locations, and whenever the "parent_id" is None, we know that's a root of a tree, and therefore perform a depth-first traversal. Find its children (wherever "parent_id" is equal to this id), use a stack to keep track of them and increment level every time/ sort alphabetically for all children of a node.

How would you be able to implement this creation of a tree (node+children) and traversal with a stack (while keeping track of level-to be able to add the hyphen- and sort)?

Would you directly traverse the JSON and do this process or create a separate structure implement and tree and then do it? Would appreciate some code for some of these different steps- I know how to solve it, I just am unclear on the exact implementation.

Answer 1

You can build this "tree" from your given data as follows:

locations = [
    {"id": 1, "name": "San Francisco Bay Area", "parent_id": None},
    {"id": 2, "name": "San Jose", "parent_id": 3},
    {"id": 3, "name": "South Bay", "parent_id": 1},
    {"id": 4, "name": "San Francisco", "parent_id": 1},
    {"id": 5, "name": "Manhattan", "parent_id": 6},
    {"id": 6, "name": "New York", "parent_id": None}
]

def find_children(parent, locations):
    branch = {}
    for location in locations:
        if location["parent_id"] == parent:
            children = find_children(location["id"], locations)
            branch[location["name"]] = children
    return branch

tree = find_children(None, locations)
print(tree)

Which prints

{'San Francisco Bay Area': {'San Francisco': {}, 'South Bay': {'San Jose': {}}}, 'New York': {'Manhattan': {}}}

You can then sort and print the content of tree :

def print_tree(tree, level=0):
    branches = sorted(list(tree.keys()))
    for branch in branches:
        print("-" * level + branch)
        print_tree(tree[branch], level + 1)

print_tree(tree)

Which prints

New York
-Manhattan
San Francisco Bay Area
-San Francisco
-South Bay
--San Jose