如何在NetworkX树中计算每个节点子节点的边权重总和

需求说明：在NetworkX构建的有向树中，为每个节点计算其所有子节点相关边的权重总和——即该节点到直接子节点的边权重，加上所有子节点各自子树中所有边的权重之和。例如：

• 非叶子节点的总和 = 自身到子节点的边权重 + 子节点的子树边权重总和
• 根节点的总和为全图所有边的权重总和
• 叶子节点（无后续子节点）的总和为0

基础示例代码 #

import networkx as nx
import matplotlib.pyplot as plt
G = nx.DiGraph()

G.add_node("ROOT")
for i in range(5):
    G.add_node("Child_%i" % i)
    G.add_node("Grandchild_%i" % i)
    G.add_node("Greatgrandchild_%i" % i)

    G.add_edge("ROOT", "Child_%i" % i, weight=(i+1))
    G.add_edge("Child_%i" % i, "Grandchild_%i" % i, weight=(2*i+1))
    G.add_edge("Grandchild_%i" % i, "Greatgrandchild_%i" % i, weight=(3*i+1))

from networkx.drawing.nx_pydot import graphviz_layout
pos = graphviz_layout(G, prog="dot")
nx.draw(G, pos, with_labels=True, arrows=True)
labels = nx.get_edge_attributes(G,'weight')
nx.draw_networkx_edge_labels(G,pos,edge_labels=labels)
plt.show()

期望计算结果示例 #

• Grandchild_4 = 13（对应边Grandchild_4→Greatgrandchild_4的权重）
• Child_4 = 22（边Child_4→Grandchild_4的权重9 + Grandchild_4的总和13）
• Grandchild_3 = 10（边Grandchild_3→Greatgrandchild_3的权重）
• Child_3 = 17（边Child_3→Grandchild_3的权重7 + Grandchild_3的总和10）
• ...
• ROOT = 全图所有边的权重总和

测试示例与方案分析 #

以下是补充的测试代码，包含两种解决方案，其中一种符合需求，另一种不符合：

import networkx as nx
import matplotlib.pyplot as plt
from pprint import pp

G = nx.DiGraph()

G.add_node("ROOT")
s = 0
for i in range(5):
    G.add_node("C_%i" % i)
    G.add_node("GC_%i" % i)
    G.add_node("GGC_%i" % i)
    G.add_node("GGD_%i" % i)

    G.add_edge("ROOT", "C_%i" % i, weight=(i+1))
    s += i+1
    G.add_edge("C_%i" % i, "GC_%i" % i, weight=(2*i+1))
    s += 2*i+1
    G.add_edge("GC_%i" % i, "GGC_%i" % i, weight=(3*i+1))
    s += 3*i+1
    G.add_edge("GC_%i" % i, "GGD_%i" % i, weight=(4 * i + 1))
    s +=4 * i + 1

print("全图边权重总和：", s)

# 方案1：基于叶子节点回溯计算（符合需求）
leaves = {n for n in G.nodes() if G.out_degree(n)==0}
out = {}
for n in leaves:
    # 获取从ROOT到当前叶子的路径边
    path_edges = next(nx.all_simple_edge_paths(G, 'ROOT', n))
    # 反向遍历路径边，累加计算
    for parent, child in reversed(path_edges):
        out[parent] = out.get(parent, 0) + out.get(child, 0) + G[parent][child]['weight']

print("\n方案1计算结果：")
pp(out)

# 方案2：基于BFS边的求和（不符合需求）
data = {
    node: sum(G.edges[edge]["weight"] for edge in nx.bfs_edges(G, node))
    for node in G.nodes
}

print("\n方案2计算结果：")
pp(data)

方案对比 #

1. 方案1（可行）：
   通过定位所有叶子节点（出度为0的节点），回溯每个叶子到根节点的路径，反向遍历路径中的边，累加当前边权重与子节点已计算的总和，最终得到父节点的子树边权重总和。利用了有向树路径唯一性的特性，结果符合需求。

2. 方案2（不可行）：
   nx.bfs_edges(G, node)仅返回当前节点到所有直接/间接子节点的边，仅统计这些边的权重之和，没有累加子节点自身的子树边总和，因此得到的结果只是当前节点到子节点的直接边权重之和，并非完整的子树边权重总和。

更简洁的通用实现 #

使用递归后序遍历的方式，先计算所有子节点的子树总和，再加上当前节点到子节点的边权重，代码简洁且通用：

def calculate_subtree_edge_sum(G, node):
    total = 0
    # 遍历当前节点的所有直接后继（子节点）
    for child in G.successors(node):
        edge_weight = G[node][child]['weight']
        # 累加当前边权重 + 子节点的子树总和
        total += edge_weight + calculate_subtree_edge_sum(G, child)
    return total

# 计算所有节点的子树边权重总和
result = {node: calculate_subtree_edge_sum(G, node) for node in G.nodes}

print("\n递归法计算结果：")
pp(result)

该方法适用于任何有向树结构，计算逻辑清晰，结果完全符合需求。

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