如何在Pandas DataFrame中通过列实现递归向量化累积权重计算

编辑说明 #

因示例数据存在错误，已移除第5行数据。

需求背景 #

现有一个由顶点V和边E构成的有向图G=(V,E)，用Pandas DataFrame存储节点间的连接关系及边的权重，示例数据如下：

#   from   to   weight
-----------------------
0     0     1     1.0
1     1     2     0.5
2     2     3     0.2     
3     0     4     1.3
4     4     5     0.9  

需要给这个DataFrame新增一列累积权重：比如路径0->1->2->3中，第2行（对应边2->3）的累积权重为1.7=0.2+0.5+1.0，最终目标DataFrame如下：

#   from   to   weight    accumulated
-------------------------------------
0     0     1     1.0      1.0
1     1     2     0.5      1.5
2     2     3     0.2      1.7   
3     0     4     1.3      1.3
4     4     5     0.9      2.2

已知前提：每个顶点仅有一条最短路径（DataFrame仅包含最短路径）。目前已有一段用DataFrame.apply实现的非向量化代码，通过字典accum_map缓存计算结果，但扩展性不足，代码如下（已修正原代码中的笔误）：

def __set_accum(self, row):
    search = row["to"]
    if search in self.accum_map:
        return self.accum_map[search]
    from_node = row["from"]
    old_from = self.df[self.df["to"] == from_node].get("from")
    old_from = None if old_from.empty else old_from.values[0]
    weight = row["weight"]
    self.accum_map[search] = self.__set_accum({"to": from_node, "from": old_from}) + weight
    return self.accum_map[search]

def set_accumulated(self):
    self.df["accumulated"] = self.df.apply(func=self.__set_accum, axis=1)

向量化实现方案 #

由于图是森林结构（多棵树，根节点为入度0的节点，比如示例中的0），我们可以通过拓扑排序批量处理节点，避免逐行操作，提升扩展性：

import pandas as pd

def add_accumulated_weight(df):
    # 构建节点到前驱节点、边权重的映射
    node_info = df.set_index('to')[['from', 'weight']].to_dict('index')
    # 计算所有节点的入度
    all_nodes = pd.concat([df['from'], df['to']]).unique()
    in_degree = df['to'].value_counts().reindex(all_nodes, fill_value=0)
    
    # 找出根节点（入度为0，是路径起点）
    root_nodes = in_degree[in_degree == 0].tolist()
    accum_map = {}
    
    # 先处理所有从根节点出发的直接边
    root_edges = df[df['from'].isin(root_nodes)]
    accum_map.update(dict(zip(root_edges['to'], root_edges['weight'])))
    processed_nodes = set(root_edges['to'])
    
    # 拓扑排序迭代处理剩余节点
    while len(processed_nodes) < len(all_nodes) - len(root_nodes):
        # 筛选出前驱节点已处理的边
        pending_edges = df[df['from'].isin(processed_nodes) & ~df['to'].isin(processed_nodes)]
        if pending_edges.empty:
            break
        
        # 批量计算累积权重：前驱节点的累积值 + 当前边权重
        pending_edges['accumulated'] = pending_edges['from'].map(accum_map) + pending_edges['weight']
        # 更新累积权重映射与已处理节点集合
        accum_map.update(dict(zip(pending_edges['to'], pending_edges['accumulated'])))
        processed_nodes.update(pending_edges['to'])
    
    # 将累积权重映射回原DataFrame
    df['accumulated'] = df['to'].map(accum_map)
    return df

# 测试示例
if __name__ == '__main__':
    sample_data = {
        'from': [0, 1, 2, 0, 4],
        'to': [1, 2, 3, 4, 5],
        'weight': [1.0, 0.5, 0.2, 1.3, 0.9]
    }
    df = pd.DataFrame(sample_data)
    result = add_accumulated_weight(df)
    print(result)

方案优势 #

• 避免了apply逐行处理的性能瓶颈，批量操作更适合大数据量场景
• 基于拓扑排序的逻辑符合图结构的遍历规则，保证计算顺序正确
• 利用Pandas的向量运算特性，整体执行效率优于递归/逐行实现

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