优化百万行DataFrame的树路径匹配遍历算法

问题背景 #

现有一个NetworkX构建的有向树，节点与Pandas DataFrame的列存在映射关系；需为百万行DataFrame的每行确定其对应的树路径。原DFS递归方案通过为每个路径新增匹配列，再从叶子列推导路径，在百万行数据下内存占用高、运行效率极低。

原方案核心问题 #

1. 为每个树路径创建单独的布尔列，百万行数据下会生成大量内存密集型临时列，导致内存占用飙升
2. 循环遍历叶子列更新path字段，多次修改DataFrame触发多次数据复制，效率低下

优化方案 #

思路核心 #

利用树的层级结构与节点-列映射关系，直接从DataFrame的每行数据中拼接出有效路径，完全避免创建临时匹配列。

步骤实现 #

1. 预处理：提取树的层级与节点父映射（可选，用于验证路径有效性）

从NetworkX树中提取每个节点的父节点及对应层级：

# 提取节点到父节点的映射
parent_map = {child: parent for parent, child in g.edges()}
# 提取每个节点的层级（根节点层级为0）
node_level = {}
def assign_level(node, level):
    node_level[node] = level
    for child in g.successors(node):
        assign_level(child, level + 1)
assign_level(root_node_id, 0)

2. 按层级排序DataFrame的列

根据节点对应的列，将列按树的层级从根到叶子排序：

# 构建列到层级的映射（同一列对应多个节点时取最小层级）
col_to_level = {}
for node, col in id_to_col_map.items():
    if col not in col_to_level or node_level[node] < col_to_level[col]:
        col_to_level[col] = node_level[node]
# 按层级排序列，得到从根到叶子的列顺序
sorted_cols = sorted(col_to_level.keys(), key=lambda x: col_to_level[x])

3. 直接拼接每行的路径

利用排序后的列，直接将每行的列值用逗号连接得到对应路径：

# 拼接路径，若存在空值可添加dropna()处理
df['path'] = df[sorted_cols].apply(lambda row: ','.join(row.astype(str)), axis=1)

4. 可选：验证路径有效性（确保每行路径在树中存在）

预先生成所有有效叶子路径，再过滤无效行：

# 生成所有从根到叶子的有效路径
valid_paths = []
def dfs_path(node, current_path):
    new_path = current_path + [node]
    if not list(g.successors(node)):
        valid_paths.append(','.join(new_path))
        return
    for child in g.successors(node):
        dfs_path(child, new_path)
dfs_path(root_node_id, [])

# 过滤无效路径
df = df[df['path'].isin(valid_paths)]

性能对比 #

• 内存占用：仅需存储最终的path列，避免了原方案中数十个临时布尔列的内存消耗
• 运行速度：直接按列拼接的时间复杂度为O(nm)（n为行数，m为层级数），远低于原方案的O(np)（p为路径数量），百万行数据下可实现秒级处理

完整优化代码 #

import networkx as nx
import pandas as pd
import numpy as np

# 初始化树与DataFrame
g = nx.DiGraph()
g.add_edge('a', 'b')
g.add_edge('a', 'c')
g.add_edge('b', 'd')
g.add_edge('b', 'e')
g.add_edge('c', 'f')
g.add_edge('c', 'g')

df = pd.DataFrame({
    "l1": "a",
    "l2": ['b', 'c', 'b', 'c'],
    'l3': ['d', 'f', 'e', 'g'],
    'data': [1, 2, 3, 4]
})

id_to_col_map = {'a': 'l1', 'b': 'l2', 'c': 'l2', 'd': 'l3', 'e': 'l3', 'f': 'l3', 'g': 'l3'}
root_node_id = 'a'

# 步骤1：提取节点父映射与层级
parent_map = {child: parent for parent, child in g.edges()}
node_level = {}
def assign_level(node, level):
    node_level[node] = level
    for child in g.successors(node):
        assign_level(child, level + 1)
assign_level(root_node_id, 0)

# 步骤2：按层级排序列
col_to_level = {}
for node, col in id_to_col_map.items():
    if col not in col_to_level or node_level[node] < col_to_level[col]:
        col_to_level[col] = node_level[node]
sorted_cols = sorted(col_to_level.keys(), key=lambda x: col_to_level[x])

# 步骤3：拼接路径
df['path'] = df[sorted_cols].apply(lambda row: ','.join(row.astype(str)), axis=1)

# 可选：验证路径有效性
valid_paths = []
def dfs_path(node, current_path):
    new_path = current_path + [node]
    if not list(g.successors(node)):
        valid_paths.append(','.join(new_path))
        return
    for child in g.successors(node):
        dfs_path(child, new_path)
dfs_path(root_node_id, [])

df = df[df['path'].isin(valid_paths)]

print(df)

内容的提问来源于stack exchange，提问作者Sreekar Mouli