Repository: scikit-learn/scikit-learn
File: sklearn/cluster/_kmeans.py (Python)

Code before:
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        n_steps = (self.max_iter * n_samples) // self._batch_size
        normalized_sample_weight = sample_weight / np.sum(sample_weight)
        unit_sample_weight = np.ones_like(sample_weight, shape=(self._batch_size,))
        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=normalized_sample_weight,
                    replace=True,
                )

                # Perform the actual update step on the minibatch data
                # Note: since the sampling of the minibatch is sample_weight aware,
                # we pass fixed unit weights to the `_mini_batch_step` call to avoid
                # accounting for the weights twice. Also note that `_mini_batch_step`
                # can be called with non-unit weights when the caller constructs
                # the batches explicitly by calling the public `partial_fit` method
                # instead.
                batch_inertia = _mini_batch_step(
                    X=X[minibatch_indices],
                    sample_weight=unit_sample_weight,
                    centers=centers,
                    centers_new=centers_new,
                    weight_sums=self._counts,
                    random_state=random_state,
                    random_reassign=self._random_reassign(),
                    reassignment_ratio=self.reassignment_ratio,
                    verbose=self.verbose,
                    n_threads=self._n_threads,
                )

                if self._tol > 0.0:
                    centers_squared_diff = np.sum((centers_new - centers) ** 2)
                else:
                    centers_squared_diff = 0

Review comment (performance):
```suggestion
        n_steps = int(self.max_iter * sample_weight.sum()) // self._batch_size
```

Code after:
                n_threads=self._n_threads,
            )

            if self.verbose:
                print(f"Inertia for init {init_idx + 1}/{self._n_init}: {inertia}")
            if best_inertia is None or inertia < best_inertia:
                init_centers = cluster_centers
                best_inertia = inertia

        centers = init_centers
        centers_new = np.empty_like(centers)

        # Initialize counts
        self._counts = np.zeros(self.n_clusters, dtype=X.dtype)

        # Attributes to monitor the convergence
        self._ewa_inertia = None
        self._ewa_inertia_min = None
        self._no_improvement = 0

        # Initialize number of samples seen since last reassignment
        self._n_since_last_reassign = 0

        sum_of_weights = np.sum(sample_weight)

        n_steps = (self.max_iter * n_samples) // self._batch_size
        normalized_sample_weight = sample_weight / sum_of_weights
        unit_sample_weight = np.ones_like(sample_weight, shape=(self._batch_size,))

        with _get_threadpool_controller().limit(limits=1, user_api="blas"):
            # Perform the iterative optimization until convergence
            for i in range(n_steps):
                # Sample a minibatch from the full dataset
                minibatch_indices = random_state.choice(
                    n_samples,
                    self._batch_size,
                    p=normalized_sample_weight,
                    replace=True,
                )
                # Perform the actual update step on the minibatch data
                # Note: since the sampling of the minibatch is sample_weight aware,
                # we pass fixed unit weights to the `_mini_batch_step` call to avoid
                # accounting for the weights twice. Also note that `_mini_batch_step`
                # can be called with non-unit weights when the caller constructs
                # the batches explicitly by calling the public `partial_fit` method
                # instead.
                batch_inertia = _mini_batch_step(
                    X=X[minibatch_indices],
                    sample_weight=unit_sample_weight,
                    centers=centers,
                    centers_new=centers_new,