ENH/FIX: in graph shortest path

Fix doctest and minor simplification
tanish2k · Feb 25, 2012 · 4395694 · 4395694
1 parent 51f5680
commit 4395694
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/sklearn/utils/graph.py b/sklearn/utils/graph.py
@@ -45,7 +45,7 @@ def single_source_shortest_path_length(graph, source, cutoff=None):
     >>> single_source_shortest_path_length(graph, 0)
     {0: 0, 1: 1, 2: 2, 3: 3}
     >>> single_source_shortest_path_length(np.ones((6, 6)), 2)
-    {2: 0, 3: 1, 4: 1, 5: 1}
+    {0: 1, 1: 1, 2: 0, 3: 1, 4: 1, 5: 1}
     """
     if sparse.isspmatrix(graph):
         graph = graph.tolil()
@@ -60,8 +60,7 @@ def single_source_shortest_path_length(graph, source, cutoff=None):
         for v in this_level:
             if v not in seen:
                 seen[v] = level     # set the level of vertex v
-                neighbors = np.array(graph.rows[v])
-                next_level.update(neighbors)
+                next_level.update(graph.rows[v])
         if cutoff is not None and cutoff <= level:
             break
         level += 1