Grouped Dataframe to a nested tree in python

Question

I have a dataframe as

grpdata = {'Group1':['A', 'A', 'A', 'B','B'],
        'Group2':['A2','B2','B2','A2','B2'],
        'Group3':['A3', 'A3', 'B3','A3', 'A3'],
        'Count':['10', '12', '14', '20']}
 
# Convert the dictionary into DataFrame 
groupdf = pd.DataFrame(grpdata)

I want to convert this dataframe to a tree, wherein each row is a path from root node to a leaf node.

I have tried using the approach shown in Read data from a pandas dataframe and create a dataframe using anytree in python

def add_nodes(nodes, parent, child):
    if parent not in nodes:
        nodes[parent] = Node(parent)  
    if child not in nodes:
        nodes[child] = Node(child)
    nodes[child].parent = nodes[parent]

nodes = {}  
for parent, child in zip(groupdf["Group1"],groupdf["Group2"]):

    add_nodes(nodes, parent, child)

However I am not able to figure out how to add the Group3 as a child to Group2 as parent node in the same node structure defined above.

Also

roots = list(groupdf[~groupdf["Group1"].isin(groupdf["Group2"])]["Group1"].unique())



for root in roots:         
    for pre, _, node in RenderTree(nodes[root]):
    print("%s%s" % (pre, node.name))

How to add the subsequent columns "Group3" and "Count to this tree structure?

score 0 · Answer 1 · answered Dec 01 '22 at 17:34

bigtree is a Python tree implementation that integrates with Python lists, dictionaries, and pandas DataFrame.

For this scenario, there is a built-in dataframe_to_tree method which does this for you.

import pandas as pd
from bigtree import dataframe_to_tree, print_tree

# I changed the dataframe to path column and count column
# I also added a `root` over here since trees must start from same root
path_data = pd.DataFrame([
    ["root/A/A2/A3", 10],
    ["root/A/B2/A3", 12],
    ["root/A/B2/B3", 14],
    ["root/B/A2/A3", 20],
    ["root/B/B2/A3", 25],
],
    columns=["Path", "Count"]
)
root = dataframe_to_tree(path_data)
print_tree(root, attr_list=["Count"])

This results in output,

root
├── A
│   ├── A2
│   │   └── A3 [Count=10]
│   └── B2
│       ├── A3 [Count=12]
│       └── B3 [Count=14]
└── B
    ├── A2
    │   └── A3 [Count=20]
    └── B2
        └── A3 [Count=25]

Source/Disclaimer: I'm the creator of bigtree ;)

Grouped Dataframe to a nested tree in python

1 Answers1