ENH : added sort_neighbors to all functions in depth_first_search.py

networkx/algorithms/traversal/breadth_first_search.py

@@ -1,5 +1,4 @@
 """Basic algorithms for breadth-first searching the nodes of a graph."""
-import math
 from collections import deque
 
 import networkx as nx
@@ -44,17 +43,17 @@ def generic_bfs_edges(G, source, neighbors=None, depth_limit=None, sort_neighbor
     depth_limit : int, optional(default=len(G))
         Specify the maximum search depth.
 
-    sort_neighbors : Callable
+    sort_neighbors : Callable (default=None)
 
         .. deprecated:: 3.2
 
            The sort_neighbors parameter is deprecated and will be removed in
            version 3.4. A custom (e.g. sorted) ordering of neighbors can be
            specified with the `neighbors` parameter.
 
-        A function that takes the list of neighbors of a given node as input,
-        and returns an iterator over these neighbors but with a custom
-        ordering.
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
 
     Yields
     ------
@@ -149,9 +148,10 @@ def bfs_edges(G, source, reverse=False, depth_limit=None, sort_neighbors=None):
     depth_limit : int, optional(default=len(G))
         Specify the maximum search depth
 
-    sort_neighbors : function
-        A function that takes the list of neighbors of given node as input, and
-        returns an *iterator* over these neighbors but with custom ordering.
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
 
     Yields
     ------
@@ -210,7 +210,7 @@ def bfs_edges(G, source, reverse=False, depth_limit=None, sort_neighbors=None):
     else:
         successors = G.neighbors
 
-    if callable(sort_neighbors):
+    if sort_neighbors is not None:
         yield from generic_bfs_edges(
             G, source, lambda node: iter(sort_neighbors(successors(node))), depth_limit
         )
@@ -236,9 +236,10 @@ def bfs_tree(G, source, reverse=False, depth_limit=None, sort_neighbors=None):
     depth_limit : int, optional(default=len(G))
         Specify the maximum search depth
 
-    sort_neighbors : function
-        A function that takes the list of neighbors of given node as input, and
-        returns an *iterator* over these neighbors but with custom ordering.
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
 
     Returns
     -------
@@ -299,9 +300,10 @@ def bfs_predecessors(G, source, depth_limit=None, sort_neighbors=None):
     depth_limit : int, optional(default=len(G))
         Specify the maximum search depth
 
-    sort_neighbors : function
-        A function that takes the list of neighbors of given node as input, and
-        returns an *iterator* over these neighbors but with custom ordering.
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
 
     Returns
     -------
@@ -364,9 +366,10 @@ def bfs_successors(G, source, depth_limit=None, sort_neighbors=None):
     depth_limit : int, optional(default=len(G))
         Specify the maximum search depth
 
-    sort_neighbors : function
-        A function that takes the list of neighbors of given node as input, and
-        returns an *iterator* over these neighbors but with custom ordering.
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
 
     Returns
     -------

networkx/algorithms/traversal/depth_first_search.py

@@ -15,7 +15,7 @@
 
 
 @nx._dispatchable
-def dfs_edges(G, source=None, depth_limit=None):
+def dfs_edges(G, source=None, depth_limit=None, *, sort_neighbors=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,11 @@ def dfs_edges(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Yields
     ------
     edge: 2-tuple of nodes
@@ -78,12 +83,18 @@ def dfs_edges(G, source=None, depth_limit=None):
     if depth_limit is None:
         depth_limit = len(G)
 
+    get_children = (
+        G.neighbors
+        if sort_neighbors is None
+        else lambda n: iter(sort_neighbors(G.neighbors(n)))
+    )
+
     visited = set()
     for start in nodes:
         if start in visited:
             continue
         visited.add(start)
-        stack = [(start, iter(G[start]))]
+        stack = [(start, get_children(start))]
         depth_now = 1
         while stack:
             parent, children = stack[-1]
@@ -92,7 +103,7 @@ 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])))
+                        stack.append((child, get_children(child)))
                         depth_now += 1
                         break
             else:
@@ -101,7 +112,7 @@ def dfs_edges(G, source=None, depth_limit=None):
 
 
 @nx._dispatchable
-def dfs_tree(G, source=None, depth_limit=None):
+def dfs_tree(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Returns oriented tree constructed from a depth-first-search from source.
 
     Parameters
@@ -114,6 +125,11 @@ def dfs_tree(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     T : NetworkX DiGraph
@@ -134,20 +150,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_neighbors=sort_neighbors))
     return T
 
 
 @nx._dispatchable
-def dfs_predecessors(G, source=None, depth_limit=None):
+def dfs_predecessors(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Returns dictionary of predecessors in depth-first-search from source.
 
     Parameters
@@ -163,6 +179,11 @@ def dfs_predecessors(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     pred: dict
@@ -194,14 +215,17 @@ 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_neighbors=sort_neighbors)
+    }
 
 
 @nx._dispatchable
-def dfs_successors(G, source=None, depth_limit=None):
+def dfs_successors(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Returns dictionary of successors in depth-first-search from source.
 
     Parameters
@@ -217,6 +241,11 @@ def dfs_successors(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     succ: dict
@@ -248,17 +277,22 @@ 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_neighbors=sort_neighbors,
+    ):
         d[s].append(t)
     return dict(d)
 
 
 @nx._dispatchable
-def dfs_postorder_nodes(G, source=None, depth_limit=None):
+def dfs_postorder_nodes(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Generate nodes in a depth-first-search post-ordering starting at source.
 
     Parameters
@@ -271,6 +305,11 @@ def dfs_postorder_nodes(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     nodes: generator
@@ -302,15 +341,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_neighbors=sort_neighbors
+    )
     return (v for u, v, d in edges if d == "reverse")
 
 
 @nx._dispatchable
-def dfs_preorder_nodes(G, source=None, depth_limit=None):
+def dfs_preorder_nodes(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Generate nodes in a depth-first-search pre-ordering starting at source.
 
     Parameters
@@ -324,6 +365,11 @@ def dfs_preorder_nodes(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     nodes: generator
@@ -355,14 +401,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_neighbors=sort_neighbors
+    )
     return (v for u, v, d in edges if d == "forward")
 
 
 @nx._dispatchable
-def dfs_labeled_edges(G, source=None, depth_limit=None):
+def dfs_labeled_edges(G, source=None, depth_limit=None, *, sort_neighbors=None):
     """Iterate over edges in a depth-first-search (DFS) labeled by type.
 
     Parameters
@@ -376,6 +424,11 @@ def dfs_labeled_edges(G, source=None, depth_limit=None):
     depth_limit : int, optional (default=len(G))
        Specify the maximum search depth.
 
+    sort_neighbors : function (default=None)
+        A function that takes an iterator over nodes as the input, and
+        returns an iterable of the same nodes with a custom ordering.
+        For example, `sorted` will sort the nodes in increasing order.
+
     Returns
     -------
     edges: generator
@@ -439,13 +492,19 @@ def dfs_labeled_edges(G, source=None, depth_limit=None):
     if depth_limit is None:
         depth_limit = len(G)
 
+    get_children = (
+        G.neighbors
+        if sort_neighbors is None
+        else lambda n: iter(sort_neighbors(G.neighbors(n)))
+    )
+
     visited = set()
     for start in nodes:
         if start in visited:
             continue
         yield start, start, "forward"
         visited.add(start)
-        stack = [(start, iter(G[start]))]
+        stack = [(start, get_children(start))]
         depth_now = 1
         while stack:
             parent, children = stack[-1]
@@ -456,7 +515,7 @@ 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])))
+                        stack.append((child, iter(get_children(child))))
                         depth_now += 1
                         break
                     else: