ENH : added sort_neighbors to all functions in depth_first_search.py

networkx/algorithms/traversal/depth_first_search.py

@@ -15,7 +15,7 @@
 
 
 @nx._dispatch
-def dfs_edges(G, source=None, depth_limit=None):
+def dfs_edges(G, source=None, depth_limit=None, sort_children=None):
     """Iterate over edges in a depth-first-search (DFS).
 
     Perform a depth-first-search over the nodes of `G` and yield
@@ -33,6 +33,10 @@ def dfs_edges(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Yields
     ------
     edge: 2-tuple of nodes
@@ -83,7 +87,7 @@ def dfs_edges(G, source=None, depth_limit=None):
         if start in visited:
             continue
         visited.add(start)
-        stack = [(start, iter(G[start]))]
+        stack = [(start, iter(sort_children(G[start]) if sort_children else G[start]))]
         depth_now = 1
         while stack:
             parent, children = stack[-1]
@@ -92,7 +96,10 @@ def dfs_edges(G, source=None, depth_limit=None):
                     yield parent, child
                     visited.add(child)
                     if depth_now < depth_limit:
-                        stack.append((child, iter(G[child])))
+                        sorted_children = (
+                            sort_children(G[child]) if sort_children else G[child]
+                        )
+                        stack.append((child, iter(sorted_children)))
                         depth_now += 1
                         break
             else:
@@ -101,7 +108,7 @@ def dfs_edges(G, source=None, depth_limit=None):
 
 
 @nx._dispatch
-def dfs_tree(G, source=None, depth_limit=None):
+def dfs_tree(G, source=None, depth_limit=None, sort_children=None):
     """Returns oriented tree constructed from a depth-first-search from source.
 
     Parameters
@@ -114,6 +121,10 @@ def dfs_tree(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     T : NetworkX DiGraph
@@ -134,20 +145,20 @@ def dfs_tree(G, source=None, depth_limit=None):
     dfs_preorder_nodes
     dfs_postorder_nodes
     dfs_labeled_edges
-    edge_dfs
-    bfs_tree
+    :func:`~networkx.algorithms.traversal.edgedfs.edge_dfs`
+    :func:`~networkx.algorithms.traversal.breadth_first_search.bfs_tree`
     """
     T = nx.DiGraph()
     if source is None:
         T.add_nodes_from(G)
     else:
         T.add_node(source)
-    T.add_edges_from(dfs_edges(G, source, depth_limit))
+    T.add_edges_from(dfs_edges(G, source, depth_limit, sort_children))
     return T
 
 
 @nx._dispatch
-def dfs_predecessors(G, source=None, depth_limit=None):
+def dfs_predecessors(G, source=None, depth_limit=None, sort_children=None):
     """Returns dictionary of predecessors in depth-first-search from source.
 
     Parameters
@@ -163,6 +174,10 @@ def dfs_predecessors(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     pred: dict
@@ -194,14 +209,14 @@ def dfs_predecessors(G, source=None, depth_limit=None):
     dfs_preorder_nodes
     dfs_postorder_nodes
     dfs_labeled_edges
-    edge_dfs
-    bfs_tree
+    :func:`~networkx.algorithms.traversal.edgedfs.edge_dfs`
+    :func:`~networkx.algorithms.traversal.breadth_first_search.bfs_tree`
     """
-    return {t: s for s, t in dfs_edges(G, source, depth_limit)}
+    return {t: s for s, t in dfs_edges(G, source, depth_limit, sort_children)}
 
 
 @nx._dispatch
-def dfs_successors(G, source=None, depth_limit=None):
+def dfs_successors(G, source=None, depth_limit=None, sort_children=None):
     """Returns dictionary of successors in depth-first-search from source.
 
     Parameters
@@ -217,6 +232,10 @@ def dfs_successors(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     succ: dict
@@ -248,17 +267,19 @@ def dfs_successors(G, source=None, depth_limit=None):
     dfs_preorder_nodes
     dfs_postorder_nodes
     dfs_labeled_edges
-    edge_dfs
-    bfs_tree
+    :func:`~networkx.algorithms.traversal.edgedfs.edge_dfs`
+    :func:`~networkx.algorithms.traversal.breadth_first_search.bfs_tree`
     """
     d = defaultdict(list)
-    for s, t in dfs_edges(G, source=source, depth_limit=depth_limit):
+    for s, t in dfs_edges(
+        G, source=source, depth_limit=depth_limit, sort_children=sort_children
+    ):
         d[s].append(t)
     return dict(d)
 
 
 @nx._dispatch
-def dfs_postorder_nodes(G, source=None, depth_limit=None):
+def dfs_postorder_nodes(G, source=None, depth_limit=None, sort_children=None):
     """Generate nodes in a depth-first-search post-ordering starting at source.
 
     Parameters
@@ -271,6 +292,10 @@ def dfs_postorder_nodes(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     nodes: generator
@@ -302,15 +327,17 @@ def dfs_postorder_nodes(G, source=None, depth_limit=None):
     dfs_edges
     dfs_preorder_nodes
     dfs_labeled_edges
-    edge_dfs
-    bfs_tree
+    :func:`~networkx.algorithms.traversal.edgedfs.edge_dfs`
+    :func:`~networkx.algorithms.traversal.breadth_first_search.bfs_tree`
     """
-    edges = nx.dfs_labeled_edges(G, source=source, depth_limit=depth_limit)
+    edges = nx.dfs_labeled_edges(
+        G, source=source, depth_limit=depth_limit, sort_children=sort_children
+    )
     return (v for u, v, d in edges if d == "reverse")
 
 
 @nx._dispatch
-def dfs_preorder_nodes(G, source=None, depth_limit=None):
+def dfs_preorder_nodes(G, source=None, depth_limit=None, sort_children=None):
     """Generate nodes in a depth-first-search pre-ordering starting at source.
 
     Parameters
@@ -324,6 +351,10 @@ def dfs_preorder_nodes(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     nodes: generator
@@ -355,14 +386,16 @@ def dfs_preorder_nodes(G, source=None, depth_limit=None):
     dfs_edges
     dfs_postorder_nodes
     dfs_labeled_edges
-    bfs_edges
+    :func:`~networkx.algorithms.traversal.breadth_first_search.bfs_edges`
     """
-    edges = nx.dfs_labeled_edges(G, source=source, depth_limit=depth_limit)
+    edges = nx.dfs_labeled_edges(
+        G, source=source, depth_limit=depth_limit, sort_children=sort_children
+    )
     return (v for u, v, d in edges if d == "forward")
 
 
 @nx._dispatch
-def dfs_labeled_edges(G, source=None, depth_limit=None):
+def dfs_labeled_edges(G, source=None, depth_limit=None, sort_children=None):
     """Iterate over edges in a depth-first-search (DFS) labeled by type.
 
     Parameters
@@ -376,6 +409,10 @@ def dfs_labeled_edges(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_children : function
+        A function that takes the list of children of given node as input, and
+        returns an *iterator* over these children but with custom ordering.
+
     Returns
     -------
     edges: generator
@@ -445,7 +482,7 @@ def dfs_labeled_edges(G, source=None, depth_limit=None):
             continue
         yield start, start, "forward"
         visited.add(start)
-        stack = [(start, iter(G[start]))]
+        stack = [(start, iter(sort_children(G[start]) if sort_children else G[start]))]
         depth_now = 1
         while stack:
             parent, children = stack[-1]
@@ -456,7 +493,10 @@ def dfs_labeled_edges(G, source=None, depth_limit=None):
                     yield parent, child, "forward"
                     visited.add(child)
                     if depth_now < depth_limit:
-                        stack.append((child, iter(G[child])))
+                        sorted_children = (
+                            sort_children(G[child]) if sort_children else G[child]
+                        )
+                        stack.append((child, iter(sorted_children)))
                         depth_now += 1
                         break
                     else:

networkx/algorithms/traversal/tests/test_dfs.py

@@ -55,6 +55,15 @@ def test_dfs_edges(self):
         edges = nx.dfs_edges(self.D)
         assert list(edges) == [(0, 1), (2, 3)]
 
+    def test_dfs_edges_sorting(self):
+        G = nx.Graph()
+        G.add_edges_from([(0, 1), (1, 2), (1, 3), (2, 4), (3, 0), (0, 4)])
+        edges_asc = nx.dfs_edges(G, source=0, sort_children=sorted)
+        sorted_desc = lambda x: sorted(x, reverse=True)
+        edges_desc = nx.dfs_edges(G, source=0, sort_children=sorted_desc)
+        assert list(edges_asc) == [(0, 1), (1, 2), (2, 4), (1, 3)]
+        assert list(edges_desc) == [(0, 4), (4, 2), (2, 1), (1, 3)]
+
     def test_dfs_labeled_edges(self):
         edges = list(nx.dfs_labeled_edges(self.G, source=0))
         forward = [(u, v) for (u, v, d) in edges if d == "forward"]
@@ -80,6 +89,53 @@ def test_dfs_labeled_edges(self):
             (0, 0, "reverse"),
         ]
 
+    def test_dfs_labeled_edges_sorting(self):
+        G = nx.Graph()
+        G.add_edges_from([(0, 1), (1, 2), (1, 3), (2, 4), (3, 0), (0, 4)])
+        edges_asc = nx.dfs_labeled_edges(G, source=0, sort_children=sorted)
+        sorted_desc = lambda x: sorted(x, reverse=True)
+        edges_desc = nx.dfs_labeled_edges(G, source=0, sort_children=sorted_desc)
+        assert list(edges_asc) == [
+            (0, 0, "forward"),
+            (0, 1, "forward"),
+            (1, 0, "nontree"),
+            (1, 2, "forward"),
+            (2, 1, "nontree"),
+            (2, 4, "forward"),
+            (4, 0, "nontree"),
+            (4, 2, "nontree"),
+            (2, 4, "reverse"),
+            (1, 2, "reverse"),
+            (1, 3, "forward"),
+            (3, 0, "nontree"),
+            (3, 1, "nontree"),
+            (1, 3, "reverse"),
+            (0, 1, "reverse"),
+            (0, 3, "nontree"),
+            (0, 4, "nontree"),
+            (0, 0, "reverse"),
+        ]
+        assert list(edges_desc) == [
+            (0, 0, "forward"),
+            (0, 4, "forward"),
+            (4, 2, "forward"),
+            (2, 4, "nontree"),
+            (2, 1, "forward"),
+            (1, 3, "forward"),
+            (3, 1, "nontree"),
+            (3, 0, "nontree"),
+            (1, 3, "reverse"),
+            (1, 2, "nontree"),
+            (1, 0, "nontree"),
+            (2, 1, "reverse"),
+            (4, 2, "reverse"),
+            (4, 0, "nontree"),
+            (0, 4, "reverse"),
+            (0, 3, "nontree"),
+            (0, 1, "nontree"),
+            (0, 0, "reverse"),
+        ]
+
     def test_dfs_labeled_disconnected_edges(self):
         edges = list(nx.dfs_labeled_edges(self.D))
         forward = [(u, v) for (u, v, d) in edges if d == "forward"]