# -*- coding: utf-8 -*-
"""
Computes minimum spanning tree of a weighted graph.
"""
# Copyright (C) 2009-2010 by
# Aric Hagberg <hagberg@lanl.gov>
# Dan Schult <dschult@colgate.edu>
# Pieter Swart <swart@lanl.gov>
# Loïc Séguin-C. <loicseguin@gmail.com>
# All rights reserved.
# BSD license.
__all__ = ['kruskal_mst',
'minimum_spanning_edges',
'minimum_spanning_tree',
'prim_mst_edges', 'prim_mst']
import networkx as nx
from heapq import heappop, heappush
from itertools import count
[docs]def minimum_spanning_edges(G, weight='weight', data=True):
"""Generate edges in a minimum spanning forest of an undirected
weighted graph.
A minimum spanning tree is a subgraph of the graph (a tree)
with the minimum sum of edge weights. A spanning forest is a
union of the spanning trees for each connected component of the graph.
Parameters
----------
G : NetworkX Graph
weight : string
Edge data key to use for weight (default 'weight').
data : bool, optional
If True yield the edge data along with the edge.
Returns
-------
edges : iterator
A generator that produces edges in the minimum spanning tree.
The edges are three-tuples (u,v,w) where w is the weight.
Examples
--------
>>> G=nx.cycle_graph(4)
>>> G.add_edge(0,3,weight=2) # assign weight 2 to edge 0-3
>>> mst=nx.minimum_spanning_edges(G,data=False) # a generator of MST edges
>>> edgelist=list(mst) # make a list of the edges
>>> print(sorted(edgelist))
[(0, 1), (1, 2), (2, 3)]
Notes
-----
Uses Kruskal's algorithm.
If the graph edges do not have a weight attribute a default weight of 1
will be used.
Modified code from David Eppstein, April 2006
http://www.ics.uci.edu/~eppstein/PADS/
"""
# Modified code from David Eppstein, April 2006
# http://www.ics.uci.edu/~eppstein/PADS/
# Kruskal's algorithm: sort edges by weight, and add them one at a time.
# We use Kruskal's algorithm, first because it is very simple to
# implement once UnionFind exists, and second, because the only slow
# part (the sort) is sped up by being built in to Python.
from networkx.utils import UnionFind
if G.is_directed():
raise nx.NetworkXError(
"Mimimum spanning tree not defined for directed graphs.")
subtrees = UnionFind()
edges = sorted(G.edges(data=True), key=lambda t: t[2].get(weight, 1))
for u, v, d in edges:
if subtrees[u] != subtrees[v]:
if data:
yield (u, v, d)
else:
yield (u, v)
subtrees.union(u, v)
[docs]def minimum_spanning_tree(G, weight='weight'):
"""Return a minimum spanning tree or forest of an undirected
weighted graph.
A minimum spanning tree is a subgraph of the graph (a tree) with
the minimum sum of edge weights.
If the graph is not connected a spanning forest is constructed. A
spanning forest is a union of the spanning trees for each
connected component of the graph.
Parameters
----------
G : NetworkX Graph
weight : string
Edge data key to use for weight (default 'weight').
Returns
-------
G : NetworkX Graph
A minimum spanning tree or forest.
Examples
--------
>>> G=nx.cycle_graph(4)
>>> G.add_edge(0,3,weight=2) # assign weight 2 to edge 0-3
>>> T=nx.minimum_spanning_tree(G)
>>> print(sorted(T.edges(data=True)))
[(0, 1, {}), (1, 2, {}), (2, 3, {})]
Notes
-----
Uses Kruskal's algorithm.
If the graph edges do not have a weight attribute a default weight of 1
will be used.
"""
T = nx.Graph(nx.minimum_spanning_edges(G, weight=weight, data=True))
# Add isolated nodes
if len(T) != len(G):
T.add_nodes_from([n for n, d in G.degree().items() if d == 0])
# Add node and graph attributes as shallow copy
for n in T:
T.node[n] = G.node[n].copy()
T.graph = G.graph.copy()
return T
kruskal_mst = minimum_spanning_tree
def prim_mst_edges(G, weight='weight', data=True):
"""Generate edges in a minimum spanning forest of an undirected
weighted graph.
A minimum spanning tree is a subgraph of the graph (a tree)
with the minimum sum of edge weights. A spanning forest is a
union of the spanning trees for each connected component of the graph.
Parameters
----------
G : NetworkX Graph
weight : string
Edge data key to use for weight (default 'weight').
data : bool, optional
If True yield the edge data along with the edge.
Returns
-------
edges : iterator
A generator that produces edges in the minimum spanning tree.
The edges are three-tuples (u,v,w) where w is the weight.
Examples
--------
>>> G=nx.cycle_graph(4)
>>> G.add_edge(0,3,weight=2) # assign weight 2 to edge 0-3
>>> mst=nx.prim_mst_edges(G,data=False) # a generator of MST edges
>>> edgelist=list(mst) # make a list of the edges
>>> print(sorted(edgelist))
[(0, 1), (1, 2), (2, 3)]
Notes
-----
Uses Prim's algorithm.
If the graph edges do not have a weight attribute a default weight of 1
will be used.
"""
if G.is_directed():
raise nx.NetworkXError(
"Mimimum spanning tree not defined for directed graphs.")
push = heappush
pop = heappop
nodes = G.nodes()
c = count()
while nodes:
u = nodes.pop(0)
frontier = []
visited = [u]
for u, v in G.edges(u):
push(frontier, (G[u][v].get(weight, 1), next(c), u, v))
while frontier:
W, _, u, v = pop(frontier)
if v in visited:
continue
visited.append(v)
nodes.remove(v)
for v, w in G.edges(v):
if not w in visited:
push(frontier, (G[v][w].get(weight, 1), next(c), v, w))
if data:
yield u, v, G[u][v]
else:
yield u, v
def prim_mst(G, weight='weight'):
"""Return a minimum spanning tree or forest of an undirected
weighted graph.
A minimum spanning tree is a subgraph of the graph (a tree) with
the minimum sum of edge weights.
If the graph is not connected a spanning forest is constructed. A
spanning forest is a union of the spanning trees for each
connected component of the graph.
Parameters
----------
G : NetworkX Graph
weight : string
Edge data key to use for weight (default 'weight').
Returns
-------
G : NetworkX Graph
A minimum spanning tree or forest.
Examples
--------
>>> G=nx.cycle_graph(4)
>>> G.add_edge(0,3,weight=2) # assign weight 2 to edge 0-3
>>> T=nx.prim_mst(G)
>>> print(sorted(T.edges(data=True)))
[(0, 1, {}), (1, 2, {}), (2, 3, {})]
Notes
-----
Uses Prim's algorithm.
If the graph edges do not have a weight attribute a default weight of 1
will be used.
"""
T = nx.Graph(nx.prim_mst_edges(G, weight=weight, data=True))
# Add isolated nodes
if len(T) != len(G):
T.add_nodes_from([n for n, d in G.degree().items() if d == 0])
# Add node and graph attributes as shallow copy
for n in T:
T.node[n] = G.node[n].copy()
T.graph = G.graph.copy()
return T