tp3_solutions.py 文件源码

def try_kmeans(X):
    """ Run the K-Means algorithm on X with different values of K, and return
     the one that gives the best score.

    Args:
        X: the TF-IDF matrix where each line represents a document and each
           column represents a word, typically obtained by running
           transform_text() from the TP2.
    """
    best_k = 1
    best_score = -1

    for k in range(2, 20+1):
        model = KMeans(n_clusters=k)
        model.fit(X)
        labels = model.predict(X)
        score = silhouette_score(model.transform(X), labels)

        print(k, "->", score)
        if score > best_score:
            best_k = k
            best_score = score

    print("The best K is", best_k)
    return best_k


# Ex3