KMeans 算法

习题 14：KMeans 算法

题目：使用 Spark 机器学习库完成无监督 KMeans 聚类任务，读取 LIBSVM 格式数据文件 input/sample_kmeans_data.txt，同时评测聚类指标。

sample_kmeans_data.txt：

0 1:2.496714 2:1.861736
0 1:2.647689 2:3.523030
0 1:1.765847 2:1.765863
0 1:3.579213 2:2.767435
0 1:1.530526 2:2.542560
0 1:1.536582 2:1.534270
0 1:2.241962 2:0.086720
0 1:0.275082 2:1.437712
0 1:0.987169 2:2.314247

import org.apache.spark.ml.clustering.KMeans
import org.apache.spark.ml.evaluation.ClusteringEvaluator
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.SparkSession

object KMeansExample {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .appName("KMeansExample")
      .master("local[*]")
      .getOrCreate()

    // 1. 读取 LIBSVM 格式数据
    val dataset = spark.read.format("libsvm").load("input/sample_kmeans_data.txt")

    println("原始数据预览：")
    dataset.show(false)

    // 2. 训练 KMeans 模型（设定 k=2，可根据实际调整）
    val kmeans = new KMeans()
      .setK(2)
      .setSeed(1L)
      .setFeaturesCol("features")

    val model = kmeans.fit(dataset)

    // 3. 获取聚类中心
    println("聚类中心：")
    model.clusterCenters.foreach(println)

    // 4. 对数据进行预测
    val predictions = model.transform(dataset)

    // 5.1 轮廓系数（Silhouette Score，值越接近 1 越好）
    val evaluator = new ClusteringEvaluator()
      .setFeaturesCol("features")
      .setPredictionCol("prediction")
      .setDistanceMeasure("squaredEuclidean")
    val silhouetteScore = evaluator.evaluate(predictions)
    println(s"轮廓系数（Silhouette Score）= $silhouetteScore")

    // 5.2 簇内平方和（WCSS / Within Set Sum of Squared Errors）
    val wcss = predictions.select("features", "prediction").rdd.map { row =>
      val features = row.getAs[org.apache.spark.ml.linalg.Vector](0)
      val clusterIdx = row.getInt(1)
      val center = model.clusterCenters(clusterIdx)
      Vectors.sqdist(features, center)
    }.sum()
    println(s"簇内平方和（WCSS）= $wcss")

    println("预测结果预览：")
    predictions.select("features", "prediction").show(false)

    spark.stop()
  }
}