def fit(self, Z):
X = Z[:, '<STR_LIT:X>'] if isinstance(Z, DictRDD) else Z<EOL>check_rdd(X, (np.ndarray, sp.spmatrix))<EOL>if self.init == '<STR_LIT>':<EOL><INDENT>self._mllib_model = MLlibKMeans.train(<EOL>X.unblock(),<EOL>self.n_clusters,<EOL>maxIterations=self.max_iter,<EOL>initializationMode="<STR_LIT>")<EOL>self.cluster_centers_ = self._mllib_model.centers<EOL><DEDENT>else:<EOL><INDENT>models = X.map(lambda X: super(SparkKMeans, self).fit(X))<EOL>models = models.map(lambda model: model.cluster_centers_).collect()<EOL>return super(SparkKMeans, self).fit(np.concatenate(models))<EOL><DEDENT>
Compute k-means clustering.

        Parameters
        ----------
        Z : ArrayRDD or DictRDD containing array-like or sparse matrix
            Train data.

        Returns
        -------
        self