ENH Avoid unnecessary O(n^2) calculation in affinity propagation (#9617)

scikit-learn · Aug 27, 2017 · 7d13009 · 7d13009
1 parent d8c363f
commit 7d13009
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sklearn/cluster/affinity_propagation_.py b/sklearn/cluster/affinity_propagation_.py
@@ -158,7 +158,7 @@ def affinity_propagation(S, preference=None, convergence_iter=15, max_iter=200,
         if verbose:
             print("Did not converge")
 
-    I = np.where(np.diag(A + R) > 0)[0]
+    I = np.flatnonzero(E)
     K = I.size  # Identify exemplars
 
     if K > 0: