Return minimum spanning edge keys for multigraphs

Before, the `minimum_spanning_edges()` function returned edges without
edge keys for multigraphs. This commit enables the user to choose to
return edge keys in addition to edges in a multigraph.

1 files changed, 113 insertions, 35 deletions

diff --git a/networkx/algorithms/tree/mst.py b/networkx/algorithms/tree/mst.py
index 09cd7651..5345db99 100644
--- a/networkx/algorithms/tree/mst.py
+++ b/networkx/algorithms/tree/mst.py
@@ -23,21 +23,42 @@ import networkx as nx
 from networkx.utils import UnionFind, not_implemented_for
 
 
-def kruskal_mst_edges(G, minimum, weight='weight', data=True):
+def kruskal_mst_edges(G, minimum, weight='weight', keys=True, data=True):
     subtrees = UnionFind()
-    edges = sorted(G.edges(data=True), key=lambda t: t[2].get(weight, 1),
-                   reverse=not minimum)
-
-    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 prim_mst_edges(G, minimum, weight='weight', data=True):
+    if G.is_multigraph():
+        edges = G.edges(keys=True, data=True)
+    else:
+        edges = G.edges(data=True)
+    getweight = lambda t: t[-1].get(weight, 1)
+    edges = sorted(edges, key=getweight, reverse=not minimum)
+    is_multigraph = G.is_multigraph()
+    # Multigraphs need to handle edge keys in addition to edge data.
+    if is_multigraph:
+        for u, v, k, d in edges:
+            if subtrees[u] != subtrees[v]:
+                if keys:
+                    if data:
+                        yield (u, v, k, d)
+                    else:
+                        yield (u, v, k)
+                else:
+                    if data:
+                        yield (u, v, d)
+                    else:
+                        yield (u, v)
+                subtrees.union(u, v)
+    else:
+        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 prim_mst_edges(G, minimum, weight='weight', keys=True, data=True):
+    is_multigraph = G.is_multigraph()
     push = heappush
     pop = heappop
 
@@ -52,24 +73,44 @@ def prim_mst_edges(G, minimum, weight='weight', data=True):
         u = nodes.pop(0)
         frontier = []
         visited = [u]
-        for u, v in G.edges(u):
-            push(frontier, (G[u][v].get(weight, 1) * sign, next(c), u, v))
-
+        if is_multigraph:
+            for u, v, k, d in G.edges(u, keys=True, data=True):
+                push(frontier, (d.get(weight, 1) * sign, next(c), u, v, k))
+        else:
+            for u, v, d in G.edges(u, data=True):
+                push(frontier, (d.get(weight, 1) * sign, next(c), u, v))
         while frontier:
-            W, _, u, v = pop(frontier)
+            if is_multigraph:
+                W, _, u, v, k = pop(frontier)
+            else:
+                W, _, u, v = pop(frontier)
             if v in visited:
                 continue
             visited.append(v)
             nodes.remove(v)
-            for v, w in G.edges(v):
-                if w in visited:
-                    continue
-                push(frontier, (G[v][w].get(weight, 1) * sign, next(c), v, w))
-
-            if data:
-                yield u, v, G[u][v]
+            if is_multigraph:
+                for _, w, k2, d2 in G.edges(v, keys=True, data=True):
+                    if w in visited:
+                        continue
+                    new_weight = d2.get(weight, 1) * sign
+                    push(frontier, (new_weight, next(c), v, w, k2))
+            else:
+                for _, w, d2 in G.edges(v, data=True):
+                    if w in visited:
+                        continue
+                    new_weight = d2.get(weight, 1) * sign
+                    push(frontier, (new_weight, next(c), v, w))
+            # Multigraphs need to handle edge keys in addition to edge data.
+            if is_multigraph and keys:
+                if data:
+                    yield u, v, k, G[u][v]
+                else:
+                    yield u, v, k
             else:
-                yield u, v
+                if data:
+                    yield u, v, G[u][v]
+                else:
+                    yield u, v
 
 ALGORITHMS = {
     'kruskal': kruskal_mst_edges,
@@ -78,17 +119,19 @@ ALGORITHMS = {
 
 
 @not_implemented_for('directed')
-def _spanning_edges(G, minimum, algorithm='kruskal', weight='weight', data=True):
+def _spanning_edges(G, minimum, algorithm='kruskal', weight='weight',
+                    keys=True, data=True):
     try:
         algo = ALGORITHMS[algorithm]
     except KeyError:
         msg = '{} is not a valid choice for an algorithm.'.format(algorithm)
         raise ValueError(msg)
 
-    return algo(G, minimum=minimum, weight=weight, data=data)
+    return algo(G, minimum=minimum, weight=weight, keys=keys, data=data)
 
 
-def minimum_spanning_edges(G, algorithm='kruskal', weight='weight', data=True):
+def minimum_spanning_edges(G, algorithm='kruskal', weight='weight', keys=True,
+                           data=True):
     """Generate edges in a minimum spanning forest of an undirected
     weighted graph.
 
@@ -109,14 +152,30 @@ def minimum_spanning_edges(G, algorithm='kruskal', weight='weight', data=True):
     weight : string
        Edge data key to use for weight (default 'weight').
 
+    keys : bool
+       Whether to yield edge key in multigraphs in addition to the
+       edge. If ``G`` is not a multigraph, this is ignored.
+
     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.
+       An iterator over tuples representing edges in a minimum spanning
+       tree of ``G``.
+
+       If ``G`` is a multigraph and both ``keys`` and ``data`` are
+       ``True``, then the tuples are four-tuples of the form ``(u, v, k,
+       w)``, where ``(u, v)`` is an edge, ``k`` is the edge key
+       identifying the particular edge joining ``u`` with ``v``, and
+       ``w`` is the weight of the edge. If ``keys`` is ``True`` but
+       ``data`` is ``False``, the tuples are three-tuples of the form
+       ``(u, v, k)``.
+
+       If ``G`` is not a multigraph, the tuples are of the form ``(u, v,
+       w)`` if ``data`` is ``True`` or ``(u, v)`` if ``data`` is
+       ``False``.
 
     Examples
     --------
@@ -150,7 +209,7 @@ def minimum_spanning_edges(G, algorithm='kruskal', weight='weight', data=True):
     http://www.ics.uci.edu/~eppstein/PADS/
     """
     return _spanning_edges(G, minimum=True, algorithm=algorithm,
-                           weight=weight, data=data)
+                           weight=weight, keys=keys, data=data)
 
 
 def maximum_spanning_edges(G, algorithm='kruskal', weight='weight', data=True):
@@ -174,14 +233,30 @@ def maximum_spanning_edges(G, algorithm='kruskal', weight='weight', data=True):
     weight : string
        Edge data key to use for weight (default 'weight').
 
+    keys : bool
+       Whether to yield edge key in multigraphs in addition to the
+       edge. If ``G`` is not a multigraph, this is ignored.
+
     data : bool, optional
        If True yield the edge data along with the edge.
 
     Returns
     -------
     edges : iterator
-       A generator that produces edges in the maximum spanning tree.
-       The edges are three-tuples (u,v,w) where w is the weight.
+       An iterator over tuples representing edges in a maximum spanning
+       tree of ``G``.
+
+       If ``G`` is a multigraph and both ``keys`` and ``data`` are
+       ``True``, then the tuples are four-tuples of the form ``(u, v, k,
+       w)``, where ``(u, v)`` is an edge, ``k`` is the edge key
+       identifying the particular edge joining ``u`` with ``v``, and
+       ``w`` is the weight of the edge. If ``keys`` is ``True`` but
+       ``data`` is ``False``, the tuples are three-tuples of the form
+       ``(u, v, k)``.
+
+       If ``G`` is not a multigraph, the tuples are of the form ``(u, v,
+       w)`` if ``data`` is ``True`` or ``(u, v)`` if ``data`` is
+       ``False``.
 
     Examples
     --------
@@ -224,7 +299,10 @@ def _optimum_spanning_tree(G, algorithm, minimum, weight='weight'):
         msg = '{} is not a valid choice for an algorithm.'.format(algorithm)
         raise ValueError(msg)
 
-    edges = algo(G, minimum=minimum, weight=weight, data=True)
+    # When creating the spanning tree, we can ignore the key used to
+    # identify multigraph edges, since a tree is guaranteed to have no
+    # multiedges. This is why we use `keys=False`.
+    edges = algo(G, minimum=minimum, weight=weight, keys=False, data=True)
     T = nx.Graph(edges)
 
     # Add isolated nodes