基于NetworkX图为含完整层级数据的Pandas DataFrame添加节点深度列

给定包含层级下属关系的Pandas DataFrame（包含直接和间接下属），且已通过NetworkX生成对应的有向层级图，我们可以利用图的路径计算能力，为每条记录添加管理者到员工的节点深度列。

示例数据 #

import pandas as pd

data = pd.DataFrame({
    "manager_id": ["A", "A", "B", "A", "C", "A", "B", "F"],
    "employee_id": ["B", "C", "C", "D", "E", "E", "E", "G"]
})

解决方案 #

假设已构建好NetworkX有向图G（层级结构为无环有向图），我们可以通过nx.shortest_path_length函数计算每条记录中管理者到员工的路径长度，该长度即为节点深度（直接下属深度为1，每多一层间接关系深度加1）。

代码实现

import networkx as nx

# 定义计算深度的函数
def get_depth(row, graph):
    return nx.shortest_path_length(graph, source=row['manager_id'], target=row['employee_id'])

# 为DataFrame添加depth列
data['depth'] = data.apply(lambda row: get_depth(row, G), axis=1)

预期输出 #

manager_id employee_id  depth
0          A           B      1
1          A           C      2
2          B           C      1
3          A           D      1
4          C           E      1
5          A           E      3
6          B           E      2
7          F           G      1

说明 #

• 层级图G需准确反映直接上下级关系（而非包含所有间接关系的全量数据），这样才能通过路径计算得到正确的深度值。
• nx.shortest_path_length在无环层级图中会返回唯一的路径长度，与题目中定义的深度完全匹配。

内容的提问来源于stack exchange，提问作者Chongxi Hong