🐛 Fix no-in-degree missing neighbors bug

eliorc · Oct 3, 2020 · 7da7572 · 7da7572
1 parent db45ac6
commit 7da7572
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/node2vec/node2vec.py b/node2vec/node2vec.py
@@ -115,9 +115,6 @@ def _precompute_probabilities(self):
                 d_graph[current_node][self.PROBABILITIES_KEY][
                     source] = unnormalized_weights / unnormalized_weights.sum()
 
-                # Save neighbors
-                d_graph[current_node][self.NEIGHBORS_KEY] = d_neighbors
-
             # Calculate first_travel weights for source
             first_travel_weights = []
 
@@ -127,6 +124,9 @@ def _precompute_probabilities(self):
             first_travel_weights = np.array(first_travel_weights)
             d_graph[source][self.FIRST_TRAVEL_KEY] = first_travel_weights / first_travel_weights.sum()
 
+            # Save neighbors
+            d_graph[source][self.NEIGHBORS_KEY] = list(self.graph.neighbors(source))
+
     def _generate_walks(self) -> list:
         """
         Generates the random walks which will be used as the skip-gram input.