-
Notifications
You must be signed in to change notification settings - Fork 0
/
mst.py
253 lines (204 loc) · 7.21 KB
/
mst.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
# -*- coding: utf-8 -*-
"""
Computes minimum spanning tree of a weighted graph.
"""
# Copyright (C) 2009-2010 by
# Aric Hagberg <[email protected]>
# Dan Schult <[email protected]>
# Pieter Swart <[email protected]>
# Loïc Séguin-C. <[email protected]>
# 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
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)
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.")
nodes = G.nodes()
while nodes:
u = nodes.pop(0)
frontier = []
visited = [u]
for u, v in G.edges(u):
heappush(frontier, (G[u][v].get(weight, 1), u, v))
while frontier:
W, u, v = heappop(frontier)
if v in visited:
continue
visited.append(v)
nodes.remove(v)
for v, w in G.edges(v):
if not w in visited:
heappush(frontier, (G[v][w].get(weight, 1), 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