scala 淘宝回头客 朴素贝叶斯代码

以下是使用Scala语言和Spark MLlib库实现淘宝回头客预测的朴素贝叶斯分类器的代码：

import org.apache.spark.ml.classification.NaiveBayes
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorAssembler}
import org.apache.spark.sql.SparkSession

// 创建SparkSession对象
val spark = SparkSession.builder()
  .appName("NaiveBayesExample")
  .getOrCreate()

// 读取数据集
val data = spark.read.format("csv")
  .option("header", "true")
  .option("inferSchema", "true")
  .load("data.csv")

// 特征选择和提取
val assembler = new VectorAssembler()
  .setInputCols(Array("category", "time_interval", "price"))
  .setOutputCol("features")
val dataset = assembler.transform(data)

// 将类别转换为数字标签
val indexer = new StringIndexer()
  .setInputCol("label")
  .setOutputCol("labelIndex")
val indexed = indexer.fit(dataset).transform(dataset)

// 划分数据集为训练集和测试集
val Array(trainingData, testData) = indexed.randomSplit(Array(0.7, 0.3), seed = 1234L)

// 训练朴素贝叶斯分类器
val nb = new NaiveBayes()
  .setSmoothing(1.0)
  .setModelType("multinomial")
val model = nb.fit(trainingData)

// 预测测试集并计算准确率
val predictions = model.transform(testData)

val evaluator = new MulticlassClassificationEvaluator()
  .setLabelCol("labelIndex")
  .setPredictionCol("prediction")
  .setMetricName("accuracy")
val accuracy = evaluator.evaluate(predictions)
println(s"Test set accuracy = $accuracy")

// 将数字标签转换为类别
val labelConverter = new IndexToString()
  .setInputCol("labelIndex")
  .setOutputCol("predictedLabel")
  .setLabels(indexer.fit(dataset).labels)
val predictionsLabel = labelConverter.transform(predictions)

// 打印预测结果
predictionsLabel.select("features", "label", "predictedLabel").show()

以上代码中，我们首先读取数据集，并对数据集进行特征选择和提取，将类别转换为数字标签。然后，我们将数据集划分为训练集和测试集，并使用训练集训练朴素贝叶斯分类器。接着，我们使用测试集进行预测，并计算预测的准确率。最后，我们将数字标签转换为类别，并打印预测结果。