Merge pull request #1756 from MisterSheik/topo-fixed

rewrite topolgical_sort, add lexicographical_topological_sort

6 files changed, 208 insertions, 152 deletions

diff --git a/doc/source/reference/algorithms.dag.rst b/doc/source/reference/algorithms.dag.rst
index bc20601d..eb006c08 100644
--- a/doc/source/reference/algorithms.dag.rst
+++ b/doc/source/reference/algorithms.dag.rst
@@ -9,7 +9,7 @@ Directed Acyclic Graphs
    ancestors
    descendants
    topological_sort
-   topological_sort_recursive
+   lexicographical_topological_sort
    is_directed_acyclic_graph
    is_aperiodic
    transitive_closure
diff --git a/networkx/algorithms/components/semiconnected.py b/networkx/algorithms/components/semiconnected.py
index ff396114..07cc05dd 100644
--- a/networkx/algorithms/components/semiconnected.py
+++ b/networkx/algorithms/components/semiconnected.py
@@ -9,7 +9,7 @@ __author__ = """ysitu <ysitu@users.noreply.github.com>"""
 # BSD license.
 
 import networkx as nx
-from networkx.utils import not_implemented_for
+from networkx.utils import not_implemented_for, pairwise
 
 __all__ = ['is_semiconnected']
 
@@ -61,4 +61,4 @@ def is_semiconnected(G):
 
     G = nx.condensation(G)
     path = nx.topological_sort(G)
-    return all(G.has_edge(u, v) for u, v in zip(path[:-1], path[1:]))
+    return all(G.has_edge(u, v) for u, v in pairwise(path))
diff --git a/networkx/algorithms/dag.py b/networkx/algorithms/dag.py
index c575c921..ae8581f6 100644
--- a/networkx/algorithms/dag.py
+++ b/networkx/algorithms/dag.py
@@ -1,24 +1,26 @@
 # -*- coding: utf-8 -*-
 """Algorithms for directed acyclic graphs (DAGs)."""
-#    Copyright (C) 2006-2011 by
+#    Copyright (C) 2006-2015 by
 #    Aric Hagberg <hagberg@lanl.gov>
 #    Dan Schult <dschult@colgate.edu>
 #    Pieter Swart <swart@lanl.gov>
 #    All rights reserved.
 #    BSD license.
 from fractions import gcd
+import heapq
 
 import networkx as nx
+from networkx.utils import consume, arbitrary_element
 from networkx.utils.decorators import *
-from ..utils import arbitrary_element
 
 __author__ = """\n""".join(['Aric Hagberg <aric.hagberg@gmail.com>',
                             'Dan Schult (dschult@colgate.edu)',
-                            'Ben Edwards (bedwards@cs.unm.edu)'])
+                            'Ben Edwards (bedwards@cs.unm.edu)',
+                            'Neil Girdhar (neil.girdhar@mcgill.ca)'])
 __all__ = ['descendants',
            'ancestors',
            'topological_sort',
-           'topological_sort_recursive',
+           'lexicographical_topological_sort',
            'is_directed_acyclic_graph',
            'is_aperiodic',
            'transitive_closure',
@@ -66,7 +68,7 @@ def ancestors(G, source):
 
 
 def is_directed_acyclic_graph(G):
-    """Return True if the graph G is a directed acyclic graph (DAG) or 
+    """Return True if the graph G is a directed acyclic graph (DAG) or
     False if not.
 
     Parameters
@@ -82,172 +84,183 @@ def is_directed_acyclic_graph(G):
     if not G.is_directed():
         return False
     try:
-        topological_sort(G, reverse=True)
+        consume(topological_sort(G))
         return True
     except nx.NetworkXUnfeasible:
         return False
 
 
-def topological_sort(G, nbunch=None, reverse=False):
-    """Return a list of nodes in topological sort order.
+def topological_sort(G):
+    """Return a generator of nodes in topologically sorted order.
 
-    A topological sort is a nonunique permutation of the nodes
-    such that an edge from u to v implies that u appears before v in the
-    topological sort order.
+    A topological sort is a nonunique permutation of the nodes such that an
+    edge from u to v implies that u appears before v in the topological sort
+    order.
 
     Parameters
     ----------
     G : NetworkX digraph
         A directed graph
 
-    nbunch : container of nodes (optional)
-        Explore graph in specified order given in nbunch
-
-    reverse : bool, optional
-        Return postorder instead of preorder if True.
-        Reverse mode is a bit more efficient.
+    Returns
+    -------
+    topologically_sorted_nodes : iterable
+        An iterable of node names in topological sorted order.
 
     Raises
     ------
     NetworkXError
-        Topological sort is defined for directed graphs only. If the
-        graph G is undirected, a NetworkXError is raised.
+        Topological sort is defined for directed graphs only. If the graph G
+        is undirected, a NetworkXError is raised.
 
     NetworkXUnfeasible
-        If G is not a directed acyclic graph (DAG) no topological sort
-        exists and a NetworkXUnfeasible exception is raised.
+        If G is not a directed acyclic graph (DAG) no topological sort exists
+        and a NetworkXUnfeasible exception is raised.  This can also be
+        raised if G is changed while the returned iterator is being processed.
+
+    RuntimeError
+        If G is changed while the returned iterator is being processed.
+
+    Examples
+    ---------
+    To get the reverse order of the topological sort::
+
+    >>> DG = nx.DiGraph([(1, 2), (2, 3)])
+    >>> list(reversed(list(nx.topological_sort(DG))))
+    [3, 2, 1]
 
     Notes
     -----
     This algorithm is based on a description and proof in
-    The Algorithm Design Manual [1]_ .
+    Introduction to algorithms - a creative approach [1]_ .
 
     See also
     --------
-    is_directed_acyclic_graph
+    is_directed_acyclic_graph, lexicographical_topological_sort
 
     References
     ----------
-    .. [1] Skiena, S. S. The Algorithm Design Manual  (Springer-Verlag, 1998). 
-        http://www.amazon.com/exec/obidos/ASIN/0387948600/ref=ase_thealgorithmrepo/
+    .. [1] Manber, U. (1989). Introduction to algorithms - a creative approach. Addison-Wesley.
+        http://www.amazon.com/Introduction-Algorithms-A-Creative-Approach/dp/0201120372
     """
     if not G.is_directed():
         raise nx.NetworkXError(
             "Topological sort not defined on undirected graphs.")
 
-    # nonrecursive version
-    seen = set()
-    order = []
-    explored = set()
-
-    if nbunch is None:
-        nbunch = G.nodes()
-    for v in nbunch:     # process all vertices in G
-        if v in explored:
-            continue
-        fringe = [v]   # nodes yet to look at
-        while fringe:
-            w = fringe[-1]  # depth first search
-            if w in explored:  # already looked down this branch
-                fringe.pop()
-                continue
-            seen.add(w)     # mark as seen
-            # Check successors for cycles and for new nodes
-            new_nodes = []
-            for n in G[w]:
-                if n not in explored:
-                    if n in seen:  # CYCLE !!
-                        raise nx.NetworkXUnfeasible("Graph contains a cycle at %s." % n)
-                    new_nodes.append(n)
-            if new_nodes:   # Add new_nodes to fringe
-                fringe.extend(new_nodes)
-            else:           # No new nodes so w is fully explored
-                explored.add(w)
-                order.append(w)
-                fringe.pop()    # done considering this node
-    if reverse:
-        return order
-    else:
-        return list(reversed(order))
+    indegree_map = {v: d for v, d in G.in_degree() if d > 0}
+    # These nodes have zero indegree and ready to be returned.
+    zero_indegree = [v for v, d in G.in_degree() if d == 0]
+
+    while zero_indegree:
+        node = zero_indegree.pop()
+        if node not in G:
+            raise RuntimeError("Graph changed during iteration")
+        for _, child in G.edges(node):
+            try:
+                indegree_map[child] -= 1
+            except KeyError:
+                raise RuntimeError("Graph changed during iteration")
+            if indegree_map[child] == 0:
+                zero_indegree.append(child)
+                del indegree_map[child]
+
+        yield node
+
+    if indegree_map:
+        raise nx.NetworkXUnfeasible("Graph contains a cycle or graph changed "
+                                    "during iteration")
 
 
-def topological_sort_recursive(G, nbunch=None, reverse=False):
-    """Return a list of nodes in topological sort order.
+def lexicographical_topological_sort(G, key=None):
+    """Return a generator of nodes in lexicographically topologically sorted
+    order.
 
-    A topological sort is a nonunique permutation of the nodes such
-    that an edge from u to v implies that u appears before v in the
-    topological sort order.
+    A topological sort is a nonunique permutation of the nodes such that an
+    edge from u to v implies that u appears before v in the topological sort
+    order.
 
     Parameters
     ----------
     G : NetworkX digraph
+        A directed graph
 
-    nbunch : container of nodes (optional)
-        Explore graph in specified order given in nbunch
+    key : function, optional
+        This function maps nodes to keys with which to resolve ambiguities in
+        the sort order.  Defaults to the identity function.
 
-    reverse : bool, optional
-        Return postorder instead of preorder if True.
-        Reverse mode is a bit more efficient.
+    Returns
+    -------
+    lexicographically_topologically_sorted_nodes : iterable
+        An iterable of node names in lexicographical topological sort order.
 
     Raises
     ------
     NetworkXError
-        Topological sort is defined for directed graphs only. If the
-        graph G is undirected, a NetworkXError is raised.
+        Topological sort is defined for directed graphs only. If the graph G
+        is undirected, a NetworkXError is raised.
 
     NetworkXUnfeasible
-        If G is not a directed acyclic graph (DAG) no topological sort
-        exists and a NetworkXUnfeasible exception is raised.
+        If G is not a directed acyclic graph (DAG) no topological sort exists
+        and a NetworkXUnfeasible exception is raised.  This can also be
+        raised if G is changed while the returned iterator is being processed.
+
+    RuntimeError
+        If G is changed while the returned iterator is being processed.
 
     Notes
     -----
-    This is a recursive version of topological sort.
+    This algorithm is based on a description and proof in
+    Introduction to algorithms - a creative approach [1]_ .
 
     See also
     --------
     topological_sort
-    is_directed_acyclic_graph
 
+    References
+    ----------
+    .. [1] Manber, U. (1989). Introduction to algorithms - a creative approach. Addison-Wesley.
+        http://www.amazon.com/Introduction-Algorithms-A-Creative-Approach/dp/0201120372
     """
     if not G.is_directed():
         raise nx.NetworkXError(
             "Topological sort not defined on undirected graphs.")
 
-    def _dfs(v):
-        ancestors.add(v)
+    if key is None:
+        key = lambda x: x
 
-        for w in G[v]:
-            if w in ancestors:
-                raise nx.NetworkXUnfeasible("Graph contains a cycle at %s." % w)
+    def create_tuple(node):
+        return key(node), node
 
-            if w not in explored:
-                _dfs(w)
+    indegree_map = {v: d for v, d in G.in_degree() if d > 0}
+    # These nodes have zero indegree and ready to be returned.
+    zero_indegree = [create_tuple(v) for v, d in G.in_degree() if d == 0]
+    heapq.heapify(zero_indegree)
 
-        ancestors.remove(v)
-        explored.add(v)
-        order.append(v)
+    while zero_indegree:
+        _, node = heapq.heappop(zero_indegree)
 
-    ancestors = set()
-    explored = set()
-    order = []
+        if node not in G:
+            raise RuntimeError("Graph changed during iteration")
+        for _, child in G.edges(node):
+            try:
+                indegree_map[child] -= 1
+            except KeyError:
+                raise RuntimeError("Graph changed during iteration")
+            if indegree_map[child] == 0:
+                heapq.heappush(zero_indegree, create_tuple(child))
+                del indegree_map[child]
 
-    if nbunch is None:
-        nbunch = G.nodes()
+        yield node
 
-    for v in nbunch:
-        if v not in explored:
-            _dfs(v)
-
-    if reverse:
-        return order
-    else:
-        return list(reversed(order))
+    if indegree_map:
+        raise nx.NetworkXUnfeasible("Graph contains a cycle or graph changed "
+                                    "during iteration")
 
 
 def is_aperiodic(G):
     """Return True if G is aperiodic.
 
-    A directed graph is aperiodic if there is no integer k > 1 that 
+    A directed graph is aperiodic if there is no integer k > 1 that
     divides the length of every cycle in the graph.
 
     Parameters
@@ -379,7 +392,7 @@ def antichains(G):
        AMS, Vol 42, 1995, p. 226.
     """
     TC = nx.transitive_closure(G)
-    antichains_stacks = [([], nx.topological_sort(G, reverse=True))]
+    antichains_stacks = [([], list(reversed(list(nx.topological_sort(G)))))]
     while antichains_stacks:
         (antichain, stack) = antichains_stacks.pop()
         # Invariant:
diff --git a/networkx/algorithms/tests/test_dag.py b/networkx/algorithms/tests/test_dag.py
index 0c784826..4fc97509 100644
--- a/networkx/algorithms/tests/test_dag.py
+++ b/networkx/algorithms/tests/test_dag.py
@@ -2,6 +2,7 @@
 from itertools import combinations
 from nose.tools import *
 from networkx.testing.utils import assert_edges_equal
+from networkx.utils import consume
 import networkx as nx
 
 
@@ -11,36 +12,28 @@ class TestDAG:
         pass
 
     def test_topological_sort1(self):
-        DG = nx.DiGraph()
-        DG.add_edges_from([(1, 2), (1, 3), (2, 3)])
-        assert_equal(nx.topological_sort(DG), [1, 2, 3])
-        assert_equal(nx.topological_sort_recursive(DG), [1, 2, 3])
+        DG = nx.DiGraph([(1, 2), (1, 3), (2, 3)])
+
+        for algorithm in [nx.topological_sort,
+                          nx.lexicographical_topological_sort]:
+            assert_equal(tuple(algorithm(DG)), (1, 2, 3))
 
         DG.add_edge(3, 2)
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort, DG)
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort_recursive, DG)
+
+        for algorithm in [nx.topological_sort,
+                          nx.lexicographical_topological_sort]:
+            assert_raises(nx.NetworkXUnfeasible, consume, algorithm(DG))
 
         DG.remove_edge(2, 3)
-        assert_equal(nx.topological_sort(DG), [1, 3, 2])
-        assert_equal(nx.topological_sort_recursive(DG), [1, 3, 2])
 
-    def test_reverse_topological_sort1(self):
-        DG = nx.DiGraph()
-        DG.add_edges_from([(1, 2), (1, 3), (2, 3)])
-        assert_equal(nx.topological_sort(DG, reverse=True), [3, 2, 1])
-        assert_equal(
-            nx.topological_sort_recursive(DG, reverse=True), [3, 2, 1])
+        for algorithm in [nx.topological_sort,
+                          nx.lexicographical_topological_sort]:
+            assert_equal(tuple(algorithm(DG)), (1, 3, 2))
 
-        DG.add_edge(3, 2)
-        assert_raises(nx.NetworkXUnfeasible,
-                      nx.topological_sort, DG, reverse=True)
-        assert_raises(nx.NetworkXUnfeasible,
-                      nx.topological_sort_recursive, DG, reverse=True)
+        DG.remove_edge(3, 2)
 
-        DG.remove_edge(2, 3)
-        assert_equal(nx.topological_sort(DG, reverse=True), [2, 3, 1])
-        assert_equal(
-            nx.topological_sort_recursive(DG, reverse=True), [2, 3, 1])
+        assert_in(tuple(nx.topological_sort(DG)), {(1, 2, 3), (1, 3, 2)})
+        assert_equal(tuple(nx.lexicographical_topological_sort(DG)), (1, 2, 3))
 
     def test_is_directed_acyclic_graph(self):
         G = nx.generators.complete_graph(2)
@@ -53,16 +46,12 @@ class TestDAG:
         DG = nx.DiGraph({1: [2], 2: [3], 3: [4],
                          4: [5], 5: [1], 11: [12],
                          12: [13], 13: [14], 14: [15]})
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort, DG)
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort_recursive, DG)
+        assert_raises(nx.NetworkXUnfeasible, consume, nx.topological_sort(DG))
 
         assert_false(nx.is_directed_acyclic_graph(DG))
 
         DG.remove_edge(1, 2)
-        assert_equal(nx.topological_sort_recursive(DG),
-                     [11, 12, 13, 14, 15, 2, 3, 4, 5, 1])
-        assert_equal(nx.topological_sort(DG),
-                     [11, 12, 13, 14, 15, 2, 3, 4, 5, 1])
+        consume(nx.topological_sort(DG))
         assert_true(nx.is_directed_acyclic_graph(DG))
 
     def test_topological_sort3(self):
@@ -77,34 +66,52 @@ class TestDAG:
             assert_equal(set(order), set(DG))
             for u, v in combinations(order, 2):
                 assert_false(nx.has_path(DG, v, u))
-        validate(nx.topological_sort_recursive(DG))
-        validate(nx.topological_sort(DG))
+        validate(list(nx.topological_sort(DG)))
 
         DG.add_edge(14, 1)
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort, DG)
-        assert_raises(nx.NetworkXUnfeasible, nx.topological_sort_recursive, DG)
+        assert_raises(nx.NetworkXUnfeasible, consume, nx.topological_sort(DG))
 
     def test_topological_sort4(self):
         G = nx.Graph()
         G.add_edge(1, 2)
-        assert_raises(nx.NetworkXError, nx.topological_sort, G)
-        assert_raises(nx.NetworkXError, nx.topological_sort_recursive, G)
+        # Only directed graphs can be topologically sorted.
+        assert_raises(nx.NetworkXError, consume, nx.topological_sort(G))
 
     def test_topological_sort5(self):
         G = nx.DiGraph()
         G.add_edge(0, 1)
-        assert_equal(nx.topological_sort_recursive(G), [0, 1])
-        assert_equal(nx.topological_sort(G), [0, 1])
-
-    def test_nbunch_argument(self):
-        G = nx.DiGraph()
-        G.add_edges_from([(1, 2), (2, 3), (1, 4), (1, 5), (2, 6)])
-        assert_equal(nx.topological_sort(G), [1, 2, 3, 6, 4, 5])
-        assert_equal(nx.topological_sort_recursive(G), [1, 5, 4, 2, 6, 3])
-        assert_equal(nx.topological_sort(G, [1]), [1, 2, 3, 6, 4, 5])
-        assert_equal(nx.topological_sort_recursive(G, [1]), [1, 5, 4, 2, 6, 3])
-        assert_equal(nx.topological_sort(G, [5]), [5])
-        assert_equal(nx.topological_sort_recursive(G, [5]), [5])
+        assert_equal(list(nx.topological_sort(G)), [0, 1])
+
+    def test_topological_sort6(self):
+        for algorithm in [nx.topological_sort,
+                          nx.lexicographical_topological_sort]:
+            def runtime_error():
+                DG = nx.DiGraph([(1, 2), (2, 3), (3, 4)])
+                first = True
+                for x in algorithm(DG):
+                    if first:
+                        first = False
+                        DG.add_edge(5 - x, 5)
+
+            def unfeasible_error():
+                DG = nx.DiGraph([(1, 2), (2, 3), (3, 4)])
+                first = True
+                for x in algorithm(DG):
+                    if first:
+                        first = False
+                        DG.remove_node(4)
+
+            def runtime_error2():
+                DG = nx.DiGraph([(1, 2), (2, 3), (3, 4)])
+                first = True
+                for x in algorithm(DG):
+                    if first:
+                        first = False
+                        DG.remove_node(2)
+
+            assert_raises(RuntimeError, runtime_error)
+            assert_raises(RuntimeError, runtime_error2)
+            assert_raises(nx.NetworkXUnfeasible, unfeasible_error)
 
     def test_ancestors(self):
         G = nx.DiGraph()
@@ -198,6 +205,28 @@ class TestDAG:
             [(1, 2, -5), (2, 3, 0), (3, 4, 1), (4, 5, 2), (3, 5, 4), (5, 6, 0), (1, 6, 2)])
         assert_equal(longest_path_length(G), 3)
 
+    def test_lexicographical_topological_sort(self):
+        G = nx.DiGraph([(1,2), (2,3), (1,4), (1,5), (2,6)])
+        assert_equal(list(nx.lexicographical_topological_sort(G)),
+                     [1, 2, 3, 4, 5, 6])
+        assert_equal(list(nx.lexicographical_topological_sort(
+            G, key=lambda x: x)),
+                     [1, 2, 3, 4, 5, 6])
+        assert_equal(list(nx.lexicographical_topological_sort(
+            G, key=lambda x: -x)),
+                     [1, 5, 4, 2, 6, 3])
+
+    def test_lexicographical_topological_sort(self):
+        G = nx.DiGraph([(1,2), (2,3), (1,4), (1,5), (2,6)])
+        assert_equal(list(nx.lexicographical_topological_sort(G)),
+                     [1, 2, 3, 4, 5, 6])
+        assert_equal(list(nx.lexicographical_topological_sort(
+            G, key=lambda x: x)),
+                     [1, 2, 3, 4, 5, 6])
+        assert_equal(list(nx.lexicographical_topological_sort(
+            G, key=lambda x: -x)),
+                     [1, 5, 4, 2, 6, 3])
+
 
 def test_is_aperiodic_cycle():
     G = nx.DiGraph()
diff --git a/networkx/relabel.py b/networkx/relabel.py
index f09cbcdc..7eb95a10 100644
--- a/networkx/relabel.py
+++ b/networkx/relabel.py
@@ -87,7 +87,7 @@ def _relabel_inplace(G, mapping):
         D = nx.DiGraph(list(mapping.items()))
         D.remove_edges_from(D.selfloop_edges())
         try:
-            nodes = nx.topological_sort(D, reverse=True)
+            nodes = reversed(list(nx.topological_sort(D)))
         except nx.NetworkXUnfeasible:
             raise nx.NetworkXUnfeasible('The node label sets are overlapping '
                                         'and no ordering can resolve the '
diff --git a/networkx/utils/misc.py b/networkx/utils/misc.py
index 9fa30922..e2084fa8 100644
--- a/networkx/utils/misc.py
+++ b/networkx/utils/misc.py
@@ -14,8 +14,10 @@ True
 #    Pieter Swart <swart@lanl.gov>
 #    All rights reserved.
 #    BSD license.
+import collections
 import sys
 import uuid
+from itertools import tee
 # itertools.accumulate is only available on Python 3.2 or later.
 #
 # Once support for Python versions less than 3.2 is dropped, this code should
@@ -175,7 +177,6 @@ def dict_to_numpy_array1(d,mapping=None):
         a[i] = d[k1]
     return a
 
-
 def is_iterator(obj):
     """Returns ``True`` if and only if the given object is an iterator
     object.
@@ -211,3 +212,16 @@ def arbitrary_element(iterable):
         raise ValueError('cannot return an arbitrary item from an iterator')
     # Another possible implementation is `for x in iterable: return x`.
     return next(iter(iterable))
+
+# Recipe from the itertools documentation.
+def consume(iterator):
+    "Consume the iterator entirely."
+    # Feed the entire iterator into a zero-length deque.
+    collections.deque(iterator, maxlen=0)
+
+# Recipe from the itertools documentation.
+def pairwise(iterable):
+    "s -> (s0, s1), (s1, s2), (s2, s3), ..."
+    a, b = tee(iterable)
+    next(b, None)
+    return zip(a, b)