Merge pull request #2 from MetrodataTeam/master

optimize fill diagnal
iskandr · Mar 15, 2017 · d045b5c · d045b5c
2 parents 7d9e650 + 0c5ba22
commit d045b5c
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/knnimpute/common.py b/knnimpute/common.py
@@ -32,6 +32,5 @@ def knn_initialize(X, missing_mask, verbose=False):
     D = all_pairs_normalized_distances(X_row_major, verbose=verbose)
     # set diagonal of distance matrix to infinity since we don't want
     # points considering themselves as neighbors
-    for i in range(X.shape[0]):
-        D[i, i] = np.inf
+    np.fill_diagonal(D, np.inf)
     return X_row_major, D