Search code examples
pythontreebreadth-first-search

python - collect full path till leaf on organization tree


I got organizations tree stored as json

{
    "name": "amos",
    "direct_reports": [
        {
            "name": "bart",
            "direct_reports": [
                {
                    "name": "colin",
                    "direct_reports": []
                },
                {
                    "name": "clara",
                    "direct_reports": []
                }
            ]
        },
        {
            "name": "bravo",
            "direct_reports": [
                {
                    "name": "cupid",
                    "direct_reports": []
                },
                {
                    "name": "clever",
                    "direct_reports": []
                }
            ]
        }
    ]
}

I need to store full "management path" for each employee, such as: management_chain["clever"]={bravo,amos} management_chain["bart"]={amos}

Currently I manage to reach all edges and classify those as employees and managers with code as followed:

def get_herarchy(org):
    tmp_obj = {}
    tmp_obj['managers'] = []
    for emp in org['direct_reports']:
        tmp_obj['managers'].append(org['name'])
        print("manager "+org['name'])
        if len(emp['direct_reports'])>0:
            get_herarchy(emp)
        tmp_obj['name'] = emp['name']
        print(emp['name'])
    return tmp_obj

But the dictionary doesn't holds the right values


Solution

  • Like this, maybe:

    def get_chain(org, name):
        if org['name'] == name:
            return [name]
        for emp in org['direct_reports']:
            chain = get_chain(emp, name)
            if chain:
                return [org['name']] + chain
        return None
    
    print(get_chain(org, 'bart'))    # ['amos', 'bart']
    print(get_chain(org, 'clever'))  # ['amos', 'bravo', 'clever']
    

    UPD: This is how to make a dictionary:

    def nested_iter(org):
        yield org['name']
        for emp in org['direct_reports']:
            yield from nested_iter(emp)
    
    print({name: get_chain(org, name)[0:-1] for name in nested_iter(org)})