Skip to content

Instantly share code, notes, and snippets.

@galbraun
galbraun / xgboost_leaf_paths_extractor.py
Created August 9, 2020 13:06
Functions to extract for a xgboost forest for each tree and each leaf - the middle nodes that create the path to reach it.
def _root_to_leaf_route(df, stack, routes):
current_node = df.loc[df.Node == stack[-1]]
if current_node.Feature.values[0] == 'Leaf':
routes[current_node.Node.values[0]] = list(stack)
stack.pop()
return
stack.append(int(current_node.Yes.values[0].split('-')[1]))
_root_to_leaf_route(df, stack, routes)
stack.append(int(current_node.No.values[0].split('-')[1]))
_root_to_leaf_route(df, stack, routes)