lb3

from sklearn.datasets import load_iris
from sklearn.cluster import MiniBatchKMeans


def load():
    X, y = load_iris(return_X_y=True)
    return X[:100]


def train_and_predict(X, n_clusters=1, random_state=42):
    model = MiniBatchKMeans(n_clusters=n_clusters, random_state=random_state)
    model.fit(X)
    preds = model.predict(X)
    return preds


def optimal_n_clusters(X, с_clusters=1, random_state=42):
    model = MiniBatchKMeans(n_clusters=с_clusters, random_state=random_state)
    model.fit(X)
    return model.inertia_