1from itertools import chain
2
3import networkx as nx
4from networkx.utils import not_implemented_for, pairwise
5
6__all__ = ["metric_closure", "steiner_tree"]
7
8
9@not_implemented_for("directed")
10@nx._dispatchable(edge_attrs="weight", returns_graph=True)
11def metric_closure(G, weight="weight"):
12 """Return the metric closure of a graph.
13
14 The metric closure of a graph *G* is the complete graph in which each edge
15 is weighted by the shortest path distance between the nodes in *G* .
16
17 Parameters
18 ----------
19 G : NetworkX graph
20
21 Returns
22 -------
23 NetworkX graph
24 Metric closure of the graph `G`.
25
26 """
27 M = nx.Graph()
28
29 Gnodes = set(G)
30
31 # check for connected graph while processing first node
32 all_paths_iter = nx.all_pairs_dijkstra(G, weight=weight)
33 u, (distance, path) = next(all_paths_iter)
34 if len(G) != len(distance):
35 msg = "G is not a connected graph. metric_closure is not defined."
36 raise nx.NetworkXError(msg)
37 Gnodes.remove(u)
38 for v in Gnodes:
39 M.add_edge(u, v, distance=distance[v], path=path[v])
40
41 # first node done -- now process the rest
42 for u, (distance, path) in all_paths_iter:
43 Gnodes.remove(u)
44 for v in Gnodes:
45 M.add_edge(u, v, distance=distance[v], path=path[v])
46
47 return M
48
49
50def _mehlhorn_steiner_tree(G, terminal_nodes, weight):
51 distances, paths = nx.multi_source_dijkstra(G, terminal_nodes, weight=weight)
52
53 d_1 = {}
54 s = {}
55 for v in G.nodes():
56 s[v] = paths[v][0]
57 d_1[(v, s[v])] = distances[v]
58
59 # G1-G4 names match those from the Mehlhorn 1988 paper.
60 G_1_prime = nx.Graph()
61 # iterate over all edges to complete d1
62 for u, v, data in G.edges(data=True):
63 su, sv = s[u], s[v]
64 weight_here = d_1[(u, su)] + data.get(weight, 1) + d_1[(v, sv)]
65 if not G_1_prime.has_edge(su, sv):
66 G_1_prime.add_edge(su, sv, weight_d1=weight_here)
67 else:
68 new_weight = min(weight_here, G_1_prime[su][sv]["weight_d1"])
69 G_1_prime.add_edge(su, sv, weight_d1=new_weight)
70
71 G_2 = nx.minimum_spanning_edges(G_1_prime, data=True, weight="weight_d1")
72
73 G_3 = nx.Graph()
74 for u, v, _ in G_2:
75 path = nx.shortest_path(G, u, v, weight=weight)
76 for n1, n2 in pairwise(path):
77 G_3.add_edge(n1, n2, weight=G[n1][n2].get(weight, 1))
78
79 G_3_mst = list(nx.minimum_spanning_edges(G_3, data=False, weight=weight))
80 if G.is_multigraph():
81 G_3_mst = (
82 (u, v, min(G[u][v], key=lambda k: G[u][v][k].get(weight, 1)))
83 for u, v in G_3_mst
84 )
85 G_4 = G.edge_subgraph(G_3_mst).copy()
86 _remove_nonterminal_leaves(G_4, terminal_nodes)
87 return G_4.edges()
88
89
90def _kou_steiner_tree(G, terminal_nodes, weight):
91 # Compute the metric closure only for terminal nodes
92 # Create a complete graph H from the metric edges
93 H = nx.Graph()
94 unvisited_terminals = set(terminal_nodes)
95
96 # check for connected graph while processing first node
97 u = unvisited_terminals.pop()
98 distances, paths = nx.single_source_dijkstra(G, source=u, weight=weight)
99 if len(G) != len(distances):
100 msg = "G is not a connected graph."
101 raise nx.NetworkXError(msg)
102 for v in unvisited_terminals:
103 H.add_edge(u, v, distance=distances[v], path=paths[v])
104
105 # first node done -- now process the rest
106 for u in unvisited_terminals.copy():
107 distances, paths = nx.single_source_dijkstra(G, source=u, weight=weight)
108 unvisited_terminals.remove(u)
109 for v in unvisited_terminals:
110 H.add_edge(u, v, distance=distances[v], path=paths[v])
111
112 # Use the 'distance' attribute of each edge provided by H.
113 mst_edges = nx.minimum_spanning_edges(H, weight="distance", data=True)
114
115 # Create an iterator over each edge in each shortest path; repeats are okay
116 mst_all_edges = chain.from_iterable(pairwise(d["path"]) for u, v, d in mst_edges)
117 if G.is_multigraph():
118 mst_all_edges = (
119 (u, v, min(G[u][v], key=lambda k: G[u][v][k].get(weight, 1)))
120 for u, v in mst_all_edges
121 )
122
123 # Find the MST again, over this new set of edges
124 G_S = G.edge_subgraph(mst_all_edges)
125 T_S = nx.minimum_spanning_edges(G_S, weight="weight", data=False)
126
127 # Leaf nodes that are not terminal might still remain; remove them here
128 T_H = G.edge_subgraph(T_S).copy()
129 _remove_nonterminal_leaves(T_H, terminal_nodes)
130
131 return T_H.edges()
132
133
134def _remove_nonterminal_leaves(G, terminals):
135 terminal_set = set(terminals)
136 leaves = {n for n in G if len(set(G[n]) - {n}) == 1}
137 nonterminal_leaves = leaves - terminal_set
138
139 while nonterminal_leaves:
140 # Removing a node may create new non-terminal leaves, so we limit
141 # search for candidate non-terminal nodes to neighbors of current
142 # non-terminal nodes
143 candidate_leaves = set.union(*(set(G[n]) for n in nonterminal_leaves))
144 candidate_leaves -= nonterminal_leaves | terminal_set
145 # Remove current set of non-terminal nodes
146 G.remove_nodes_from(nonterminal_leaves)
147 # Find any new non-terminal nodes from the set of candidates
148 leaves = {n for n in candidate_leaves if len(set(G[n]) - {n}) == 1}
149 nonterminal_leaves = leaves - terminal_set
150
151
152ALGORITHMS = {
153 "kou": _kou_steiner_tree,
154 "mehlhorn": _mehlhorn_steiner_tree,
155}
156
157
158@not_implemented_for("directed")
159@nx._dispatchable(preserve_all_attrs=True, returns_graph=True)
160def steiner_tree(G, terminal_nodes, weight="weight", method=None):
161 r"""Return an approximation to the minimum Steiner tree of a graph.
162
163 The minimum Steiner tree of `G` w.r.t a set of `terminal_nodes` (also *S*)
164 is a tree within `G` that spans those nodes and has minimum size (sum of
165 edge weights) among all such trees.
166
167 The approximation algorithm is specified with the `method` keyword
168 argument. All three available algorithms produce a tree whose weight is
169 within a ``(2 - (2 / l))`` factor of the weight of the optimal Steiner tree,
170 where ``l`` is the minimum number of leaf nodes across all possible Steiner
171 trees.
172
173 * ``"kou"`` [2]_ (runtime $O(|S| |V|^2)$) computes the minimum spanning tree of
174 the subgraph of the metric closure of *G* induced by the terminal nodes,
175 where the metric closure of *G* is the complete graph in which each edge is
176 weighted by the shortest path distance between the nodes in *G*.
177
178 * ``"mehlhorn"`` [3]_ (runtime $O(|E|+|V|\log|V|)$) modifies Kou et al.'s
179 algorithm, beginning by finding the closest terminal node for each
180 non-terminal. This data is used to create a complete graph containing only
181 the terminal nodes, in which edge is weighted with the shortest path
182 distance between them. The algorithm then proceeds in the same way as Kou
183 et al..
184
185 Parameters
186 ----------
187 G : NetworkX graph
188
189 terminal_nodes : list
190 A list of terminal nodes for which minimum steiner tree is
191 to be found.
192
193 weight : string (default = 'weight')
194 Use the edge attribute specified by this string as the edge weight.
195 Any edge attribute not present defaults to 1.
196
197 method : string, optional (default = 'mehlhorn')
198 The algorithm to use to approximate the Steiner tree.
199 Supported options: 'kou', 'mehlhorn'.
200 Other inputs produce a ValueError.
201
202 Returns
203 -------
204 NetworkX graph
205 Approximation to the minimum steiner tree of `G` induced by
206 `terminal_nodes` .
207
208 Raises
209 ------
210 NetworkXNotImplemented
211 If `G` is directed.
212
213 ValueError
214 If the specified `method` is not supported.
215
216 Notes
217 -----
218 For multigraphs, the edge between two nodes with minimum weight is the
219 edge put into the Steiner tree.
220
221
222 References
223 ----------
224 .. [1] Steiner_tree_problem on Wikipedia.
225 https://en.wikipedia.org/wiki/Steiner_tree_problem
226 .. [2] Kou, L., G. Markowsky, and L. Berman. 1981.
227 ‘A Fast Algorithm for Steiner Trees’.
228 Acta Informatica 15 (2): 141–45.
229 https://doi.org/10.1007/BF00288961.
230 .. [3] Mehlhorn, Kurt. 1988.
231 ‘A Faster Approximation Algorithm for the Steiner Problem in Graphs’.
232 Information Processing Letters 27 (3): 125–28.
233 https://doi.org/10.1016/0020-0190(88)90066-X.
234 """
235 if method is None:
236 method = "mehlhorn"
237
238 try:
239 algo = ALGORITHMS[method]
240 except KeyError as e:
241 raise ValueError(f"{method} is not a valid choice for an algorithm.") from e
242
243 edges = algo(G, terminal_nodes, weight)
244 # For multigraph we should add the minimal weight edge keys
245 if G.is_multigraph():
246 edges = (
247 (u, v, min(G[u][v], key=lambda k: G[u][v][k][weight])) for u, v in edges
248 )
249 T = G.edge_subgraph(edges)
250 return T