fix second node of weighted walk

louisabraham · Jun 21, 2020 · 948a5f2 · 948a5f2
1 parent c6ba426
commit 948a5f2
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/fastnode2vec/graph.py b/fastnode2vec/graph.py
@@ -69,7 +69,9 @@ def _random_walk_weighted(indptr, indices, data, walk_length, p, q, t):
 
     walk = np.empty(walk_length, dtype=indices.dtype)
     walk[0] = t
-    walk[1] = np.random.choice(_neighbors(indptr, indices, t))
+    walk[1] = _neighbors(indptr, indices, t)[
+        np.searchsorted(_neighbors(indptr, data, t), np.random.rand())
+    ]
     for j in range(2, walk_length):
         neighbors = _neighbors(indptr, indices, walk[j - 1])
         neighbors_p = _neighbors(indptr, data, walk[j - 1])