Search code examples
pythongraphtreenetworkx

How to get for each node the sum of the edge weights for children in a networkx tree


In the example below, how to get the sum of all the children edges from the ROOT.

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()

For example in the example below, I would like to have for every node:

  • Grandchild_4 = 13
  • Child_4 = 22
  • GrandChild_3 = 10
  • Child_3 = 17
  • ...
  • ROOT = sum of all the edges on this graph

enter image description here

EDIT: here is an example which works in the accepted solution, but not the other one:

import networkx as nx
import matplotlib.pyplot as plt
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)

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()

###
leaves = {n for n in G.nodes() if G.out_degree(n)==0}

out = {}

for n in leaves:
    for parent, child in reversed(next(nx.all_simple_edge_paths(G, 'ROOT', n))):
        out[parent] = (out.get(parent, 0) + out.get(child, 0)
                       + G.get_edge_data(parent, child)['weight']
                      )
from pprint import pp
pp(out)


data = {
    node: sum(G.edges[edge]["weight"] for edge in nx.bfs_edges(G, node))
    for node in G.nodes
}

from pprint import pp
pp(data)

Solution

  • If you want a method that can work for any starting node, you can use an iterative process using bfs_edges for example:

    total_weights = sum(G.edges[edge]["weight"] for edge in nx.bfs_edges(G, source))
    

    You then need to iterate over each of the nodes to have what you want. As a dictionary, this would be:

    data = {
        node: sum(G.edges[edge]["weight"] for edge in nx.bfs_edges(G, node))
        for node in G.nodes
    }