Identifying root parents and all their children in trees
I suggest you use networkx, as this is a graph problem. In particular the descendants function:
import networkx as nx
import pandas as pd
data = [['A', 'B', 0, 1],
['B', 'C', 1, 2],
['B', 'D', 1, 2],
['X', 'Y', 0, 2],
['X', 'D', 0, 2],
['Y', 'Z', 2, 3]]
df = pd.DataFrame(data=data, columns=['parent', 'child', 'parent_level', 'child_level'])
roots = df.parent[df.parent_level.eq(0)].unique()
dg = nx.from_pandas_edgelist(df, source='parent', target='child', create_using=nx.DiGraph)
result = pd.DataFrame(data=[[root, nx.descendants(dg, root)] for root in roots], columns=['root', 'children'])
print(result)
Output
root children
0 A {D, B, C}
1 X {Z, Y, D}
With Recursion
def find_root(tree, child):
if child in tree:
return {p for x in tree[child] for p in find_root(tree, x)}
else:
return {child}
tree = {}
for parent, child in zip(df.parent, df.child):
tree.setdefault(child, set()).add(parent)
descendents = {}
for child in tree:
for parent in find_root(tree, child):
descendents.setdefault(parent, set()).add(child)
pd.DataFrame(descendents.items(), columns=['root', 'children'])
root children
0 A {B, D, C}
1 X {Z, D, Y}
You could alternatively set up find_root
as a generator
def find_root(tree, child):
if child in tree:
for x in tree[child]:
yield from find_root(tree, x)
else:
yield child
Further, if you want to avoid recursion depth issues, you can use the "stack of iterators" pattern to define find_root
def find_root(tree, child):
stack = [iter([child])]
while stack:
for node in stack[-1]:
if node in tree:
stack.append(iter(tree[node]))
else:
yield node
break
else: # yes! that is an `else` clause on a for loop
stack.pop()