Spark决策树代码及注解

一、简介

决策树是一种简单但广泛使用的分类器，它通过训练数据构建决策树，对未知的数据进行分类。

二、示例

1、数据

以下是示例数据，该数据下载地址：机器学习文件数据包。

0 128:51 129:159 130:253 131:159 132:50 155:48 156:238 157:252 158:252 159:252 160:237 182:54 183:227 184:253 185:252 186:239 187:233 188:252 189:57 190:6 208:10 209:60 210:224 211:252 212:253 213:252 214:202 215:84 216:252 217:253 218:122 236:163 237:252 238:252 239:252 240:253 241:252 242:252 243:96 244:189 245:253 246:167 263:51 264:238 265:253 266:253 267:190 268:114 269:253 270:228 271:47 272:79 273:255 274:168 290:48 291:238 292:252 293:252 294:179 295:12 296:75 297:121 298:21 301:253 302:243 303:50 317:38 318:165 319:253 320:233 321:208 322:84 329:253 330:252 331:165 344:7 345:178 346:252 347:240 348:71 349:19 350:28 357:253 358:252 359:195 372:57 373:252 374:252 375:63 385:253 386:252 387:195 400:198 401:253 402:190 413:255 414:253 415:196 427:76 428:246 429:252 430:112 441:253 442:252 443:148 455:85 456:252 457:230 458:25 467:7 468:135 469:253 470:186 471:12 483:85 484:252 485:223 494:7 495:131 496:252 497:225 498:71 511:85 512:252 513:145 521:48 522:165 523:252 524:173 539:86 540:253 541:225 548:114 549:238 550:253 551:162 567:85 568:252 569:249 570:146 571:48 572:29 573:85 574:178 575:225 576:253 577:223 578:167 579:56 595:85 596:252 597:252 598:252 599:229 600:215 601:252 602:252 603:252 604:196 605:130 623:28 624:199 625:252 626:252 627:253 628:252 629:252 630:233 631:145 652:25 653:128 654:252 655:253 656:252 657:141 658:37
1 159:124 160:253 161:255 162:63 186:96 187:244 188:251 189:253 190:62 214:127 215:251 216:251 217:253 218:62 241:68 242:236 243:251 244:211 245:31 246:8 268:60 269:228 270:251 271:251 272:94 296:155 297:253 298:253 299:189 323:20 324:253 325:251 326:235 327:66 350:32 351:205 352:253 353:251 354:126 378:104 379:251 380:253 381:184 382:15 405:80 406:240 407:251 408:193 409:23 432:32 433:253 434:253 435:253 436:159 460:151 461:251 462:251 463:251 464:39 487:48 488:221 489:251 490:251 491:172 515:234 516:251 517:251 518:196 519:12 543:253 544:251 545:251 546:89 570:159 571:255 572:253 573:253 574:31 597:48 598:228 599:253 600:247 601:140 602:8 625:64 626:251 627:253 628:220 653:64 654:251 655:253 656:220 681:24 682:193 683:253 684:220

2、代码

package com.ml.decisionTree

import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.mllib.tree.DecisionTree
import org.apache.spark.mllib.util.MLUtils

object DecisionTreeDemo {
    def main(args: Array[String]): Unit = {
        val sc = new SparkContext(new SparkConf().setMaster("local").setAppName(this.getClass.getSimpleName))
        val file = MLUtils.loadLibSVMFile(sc, "src/main/resources/svm/sample_libsvm_data.txt")

        // 切分数据
        val array = file.randomSplit(Array(0.7, 0.3), 5)
        val (trainData, testData) = (array(0), array(1))

        // 指定分类数目
        val numClasses = 2
        // 设置分类特征字段信息，离散/分类特征字段，这里设置为空
        val categoricalFeatures = Map[Int, Int]()
        /*
        * 决策树的impurity评估方法（划分的度量选择）：gini基尼系数，entropy熵
        * 即基尼系数或信息熵，两种算法的准确率差异不大，信息熵的效率更低一些（因为它有对数运算）；
        * 一般使用默认的基尼系数”gini”即CART算法，也可以选择ID3, C4.5的最优特征选择方法。
        * */
        val impurity = "gini"
        // 树的最大深度
        val maxDepth = 5
        // 离散连续特征时使用的bin数。增加maxBins允许算法考虑更多的分割候选者并进行细粒度的分割决策。
        val maxBins = 32

        // 构建模型
        val model = DecisionTree.trainClassifier(trainData, numClasses, categoricalFeatures, impurity, maxDepth, maxBins)

        // 预测
        val predictions = testData.map { point =>
            // 根据特征预测分类
            val prediction = model.predict(point.features)
            (point.label, prediction)
        }

        // 打印预测结果
        println("正确率：" + predictions.filter(r => r._1 == r._2).count().toDouble / testData.count())
    }
}