Pandas: determine id of tree of categories from two columns

user_1571934 0
Viewed 127

I have a tree of categories setup as follows. The top level is defined by parent_id = -1 (in this case I have two nodes (i.e. Linear Asset and Point Asset) at the top level.

asset_tree = [
    {'id': 1, 'name': 'Linear Asset', 'parent_id': -1},
    {'id': 2, 'name': 'Lateral', 'parent_id': 1},
    {'id': 3, 'name': 'Main', 'parent_id': 1},
    {'id': 4, 'name': 'Point Asset', 'parent_id': -1},
    {'id': 5, 'name': 'Fountain', 'parent_id': 4},
    {'id': 6, 'name': 'Hydrant', 'parent_id': 4}
]

I also have a dataframe of assets defined as this:

import pandas as pd

df = pd.DataFrame({
    'name': ['pipe_1','pipe_2','pipe_3','hydrant_1', 'hydrant_2', 'fountain_1', 'fountain_2'],
    'level_1': ['Linear Asset','Linear Asset','Linear Asset','Point Asset','Point Asset','Point Asset','Point Asset'],
    'level_2': ['Main','Lateral','Lateral','Hydrant','Hydrant','Fountain','Fountain']
})

and thus the dataframe looks like this:

         name       level_1   level_2
0      pipe_1  Linear Asset      Main
1      pipe_2  Linear Asset   Lateral
2      pipe_3  Linear Asset   Lateral
3   hydrant_1   Point Asset   Hydrant
4   hydrant_2   Point Asset   Hydrant
5  fountain_1   Point Asset  Fountain
6  fountain_2   Point Asset  Fountain

I want a function that will lookup the id of lowest level of the tree (in the example case, level_2). For my example code, my dataframe output would be as follows. Also, I would like the function to work if I had 1, 2, or 3 levels.

         name       level_1   level_2 tree_id
0      pipe_1  Linear Asset      Main       3
1      pipe_2  Linear Asset   Lateral       2
2      pipe_3  Linear Asset   Lateral       2
3   hydrant_1   Point Asset   Hydrant       6
4   hydrant_2   Point Asset   Hydrant       6
5  fountain_1   Point Asset  Fountain       5
6  fountain_2   Point Asset  Fountain       5

I thought of the following function, but there are a few problems:

  • it does not work (it gives me the level_1 id, not the level_2)
  • ideally, if I had three levels (i.e. level_1, level_2, and level_3) instead of just two, it would work with a more complex tree.
def find_tree_id(branches, tree):
    tree_id = None
    number_of_branches = len(branches)
    parent_id = -99
    for i in range(0,number_of_branches):
        for j in range(0,len(tree)):
            if i == 0 and branches[i] == tree[j]['name']:
                parent_id = tree[j]['parent_id']
                tree_id = tree[j]['id']
            if parent_id == -1:
                return tree_id
    return tree_id

tree_ids = []
for i, row in df.iterrows():
    tree_id = find_tree_id([row['level_1'], row['level_2']], asset_tree)
    tree_ids.append(tree_id)
df['tree_id'] = tree_ids
print(df)

The erroneous output is:

         name       level_1   level_2  tree_id
0      pipe_1  Linear Asset      Main        1
1      pipe_2  Linear Asset   Lateral        1
2      pipe_3  Linear Asset   Lateral        1
3   hydrant_1   Point Asset   Hydrant        4
4   hydrant_2   Point Asset   Hydrant        4
5  fountain_1   Point Asset  Fountain        4
6  fountain_2   Point Asset  Fountain        4
pandas python
2 Answers

You can first transform asset_tree into a nested dictionary, storing the relationships between the levels. This way, you can then use a recursive generator function that takes in a level row and traverses new tree, using the names in the row to get the id of the right most name in the level:

import pandas as pd
asset_tree = [{'id': 1, 'name': 'Linear Asset', 'parent_id': -1}, {'id': 2, 'name': 'Lateral', 'parent_id': 1}, {'id': 3, 'name': 'Main', 'parent_id': 1}, {'id': 4, 'name': 'Point Asset', 'parent_id': -1}, {'id': 5, 'name': 'Fountain', 'parent_id': 4}, {'id': 6, 'name': 'Hydrant', 'parent_id': 4}]
a_tree = {i['id']:i for i in asset_tree}
assets = {'name': ['pipe_1', 'pipe_2', 'pipe_3', 'hydrant_1', 'hydrant_2', 'fountain_1', 'fountain_2'], 'level_1': ['Linear Asset', 'Linear Asset', 'Linear Asset', 'Point Asset', 'Point Asset', 'Point Asset', 'Point Asset'], 'level_2': ['Main', 'Lateral', 'Lateral', 'Hydrant', 'Hydrant', 'Fountain', 'Fountain']}
def to_dict(_id):
  return {i['id']:to_dict(i['id']) for i in asset_tree if i['parent_id'] == _id}

new_tree = {i['id']:to_dict(i['id']) for i in asset_tree if i['parent_id'] == -1}
def get_id(d, row, c = []):
   if not row:
      yield c
   else:
      for a, b in d.items():
         if row[0] == a_tree[a]['name']:
            yield from get_id(b, row[1:], c+[a])

result = [dict(zip([*assets, 'tree_id'], [a, *b, next(get_id(new_tree, b))[-1]])) 
          for a, *b in zip(*assets.values())]
df = pd.DataFrame(result)

Output:

         name       level_1   level_2  tree_id
0      pipe_1  Linear Asset      Main        3
1      pipe_2  Linear Asset   Lateral        2
2      pipe_3  Linear Asset   Lateral        2
3   hydrant_1   Point Asset   Hydrant        6
4   hydrant_2   Point Asset   Hydrant        6
5  fountain_1   Point Asset  Fountain        5
6  fountain_2   Point Asset  Fountain        5
user_73267380

Here is the solution I came up with. I believe it is more robust than the answers given to-date, but not perfect.

I turn my tree into a dataframe, and wrote a recursive function to derive the concatenated names. The recurson allows for the flexibility of have more levels if needed. I add the names as a column called 'flat_levels'.

def get_flatname_from_tree(tree_df):
    return tree_df['id'].apply(_flatname_recurse, df=tree_df)

def _flatname_recurse(ID, df):
    row = df.iloc[ID - 1]
    if row['parent_id'] == -1:
        return row['name']
    else:
        return _flatname_recurse(row['parent_id'], df=df) + ' : ' + row['name']

tree = pd.DataFrame(asset_tree)
tree['flat_levels'] = get_flatname_from_tree(tree)

then I wrote a function to concatenate the levels in my df

def concatentate_levels(x, num_levels):
    all_levels = ""
    for i in range(1,num_levels+1):
        all_levels = all_levels + x['level_' + str(i)] + " : "
    all_levels = all_levels[:-3]  # removes the trailing colon
    return all_levels

n_levels = len(df.filter(like='level_').columns)
df['flat_levels'] = df.apply(concatentate_levels, args=[n_levels], axis=1)

Finally, I merge the two dataframes on the 'flat_levels' column

df = pd.merge(df, tree[["id","flat_levels"]], on="flat_levels")
user_15719340
Related