Python实现决策树ID3算法

主要思想：

　　0、训练集格式：特征1,特征2,...特征n,类别

　　1、采用Python自带的数据结构字典递归的表示数据

　　2、ID3计算的信息增益是指类别的信息增益，因此每次都是计算类别的熵

　　3、ID3每次选择最优特征进行数据划分后都会消耗特征

　　4、当特征消耗到一定程度，可能会出现数据实例一样，但是类别不一样的情况，这个时候选不出最优特征而返回-1；

　　　  因此外面要捕获-1，要不然Python会以为最优特征是最后一列(类别)

#coding=utf-8

import operator

from math import log

import time

import os, sys

import string

def createDataSet(trainDataFile):

    print trainDataFile

    dataSet = []

    try:

        fin = open(trainDataFile)

        for line in fin:

            line = line.strip()

            cols = line.split('\t')

            row = [cols[1], cols[2], cols[3], cols[4], cols[5], cols[6], cols[7], cols[8], cols[9], cols[10], cols[0]]

            dataSet.append(row)

            #print row

    except:

        print 'Usage xxx.py trainDataFilePath outputTreeFilePath'

        sys.exit()

        labels = ['cip1', 'cip2', 'cip3', 'cip4', 'sip1', 'sip2', 'sip3', 'sip4', 'sport', 'domain']

    print 'dataSetlen', len(dataSet)

        return dataSet, labels

#calc shannon entropy

def calcShannonEnt(dataSet):

    numEntries = len(dataSet)

    labelCounts = {}

    for feaVec in dataSet:

        currentLabel = feaVec[-1]　　#每次都是计算类别的熵

        if currentLabel not in labelCounts:

            labelCounts[currentLabel] = 0

        labelCounts[currentLabel] += 1

    shannonEnt = 0.0

    for key in labelCounts:

        prob = float(labelCounts[key])/numEntries

        shannonEnt -= prob * log(prob, 2)

    return shannonEnt

def splitDataSet(dataSet, axis, value):

    retDataSet = []

    for featVec in dataSet:

        if featVec[axis] == value:

            reducedFeatVec = featVec[:axis]

            reducedFeatVec.extend(featVec[axis+1:])

            retDataSet.append(reducedFeatVec)

    return retDataSet

def chooseBestFeatureToSplit(dataSet):

    numFeatures = len(dataSet[0]) - 1    #last col is label

    baseEntropy = calcShannonEnt(dataSet)

    bestInfoGain = 0.0

    bestFeature = -1

    for i in range(numFeatures):

        featList = [example[i] for example in dataSet]

        uniqueVals = set(featList)

        newEntropy = 0.0

        for value in uniqueVals:

            subDataSet = splitDataSet(dataSet, i, value)

            prob = len(subDataSet) / float(len(dataSet))

            newEntropy += prob * calcShannonEnt(subDataSet)

        infoGain = baseEntropy -newEntropy

        if infoGain > bestInfoGain:

            bestInfoGain = infoGain

            bestFeature = i

    return bestFeature

#feature is exhaustive, reture what you want label

def majorityCnt(classList):

    classCount = {}

    for vote in classList:

        if vote not in classCount.keys():

            classCount[vote] = 0

        classCount[vote] += 1

    return max(classCount)

def createTree(dataSet, labels):

    classList = [example[-1] for example in dataSet]

    if classList.count(classList[0]) ==len(classList):    #all data is the same label

        return classList[0]

    if len(dataSet[0]) == 1:    #all feature is exhaustive

        return majorityCnt(classList)

    bestFeat = chooseBestFeatureToSplit(dataSet)

    bestFeatLabel = labels[bestFeat]

    if(bestFeat == -1):        #特征一样，但类别不一样，即类别与特征不相关，随机选第一个类别做分类结果

        return classList[0]

    myTree = {bestFeatLabel:{}}

    del(labels[bestFeat])

    featValues = [example[bestFeat] for example in dataSet]

    uniqueVals = set(featValues)

    for value in uniqueVals:

        subLabels = labels[:]

        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value),subLabels)

    return myTree

def main():

    data,label = createDataSet(sys.argv[1])

    t1 = time.clock()

    myTree = createTree(data,label)

    t2 = time.clock()

    fout = open(sys.argv[2], 'w')

    fout.write(str(myTree))

    fout.close()

    print 'execute for ',t2-t1

if __name__=='__main__':

    main()