Warning

This documents an unmaintained version of NetworkX. Please upgrade to a maintained version and see the current NetworkX documentation.

# Source code for networkx.algorithms.mst

# -*- coding: utf-8 -*-
"""
Computes minimum spanning tree of a weighted graph.

"""
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    Loïc Séguin-C. <loicseguin@gmail.com>

__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
"""
# Modified code from David Eppstein, April 2006
# 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))
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))