ENH use np.cumsum instead of stable_cumsum in kmeans++ (scikit-learn#31991)

otizonaizit · web-flow · commit d5715fb24288 · 2025-08-22T19:41:50.000+05:00
diff --git a/doc/whats_new/upcoming_changes/sklearn.cluster/31991.efficiency.rst b/doc/whats_new/upcoming_changes/sklearn.cluster/31991.efficiency.rst
@@ -0,0 +1,3 @@
+- :func:`cluster.kmeans_plusplus` now uses `np.cumsum` directly without extra
+  numerical stability checks and without casting to `np.float64`.
+  By :user:`Tiziano Zito <otizonaizit>`
diff --git a/sklearn/cluster/_kmeans.py b/sklearn/cluster/_kmeans.py
@@ -42,7 +42,7 @@
 from sklearn.utils import check_array, check_random_state
 from sklearn.utils._openmp_helpers import _openmp_effective_n_threads
 from sklearn.utils._param_validation import Interval, StrOptions, validate_params
-from sklearn.utils.extmath import row_norms, stable_cumsum
+from sklearn.utils.extmath import row_norms
 from sklearn.utils.parallel import (
     _get_threadpool_controller,
     _threadpool_controller_decorator,
@@ -248,7 +248,7 @@ def _kmeans_plusplus(
         # to the squared distance to the closest existing center
         rand_vals = random_state.uniform(size=n_local_trials) * current_pot
         candidate_ids = np.searchsorted(
-            stable_cumsum(sample_weight * closest_dist_sq), rand_vals
+            np.cumsum(sample_weight * closest_dist_sq), rand_vals
         )
         # XXX: numerical imprecision can result in a candidate_id out of range
         np.clip(candidate_ids, None, closest_dist_sq.size - 1, out=candidate_ids)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+- :func:`cluster.kmeans_plusplus` now uses `np.cumsum` directly without extra
	`2`	+ numerical stability checks and without casting to `np.float64`.
	`3`	+ By :user:`Tiziano Zito <otizonaizit>`