Source code for networkx.algorithms.mst

# -*- 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