使用pandas的dataFrame实现的决策树算法

使用pandas的dataFrame实现的决策 树算法

网上的决策树算法实现多使用的是列表的数据实现的，但是实际使用中DataFrame对数据前期处理更加友好。

因此编写了此决策树算法与大家分享。其中代码多有简陋之处，请大家留言改进。

决策树算法的基本介绍已经有很多文章了，这里不作过多介绍，请自行寻找。

伪代码：

训练创建树或子树：

如果数据集中数据记录为同一y分类，则返回此分类标签

如果数据集仅剩下2列，则返回此数据集中y分类较多的一类的标签（2列数据，1个是变量，另一个是y标签）

否则需要创建子树，调用最佳分类特征选择函数，选择最佳分类特征

确定最佳特征的分枝

for 每一分枝：

          递归调用创建树方法，创建子树。

最佳分类特征选择函数：

计算基本信息熵

for 每一特征：

计算特征分枝后的信息熵

计算特征分枝后信息增益

选择最大信息增益

返回最佳分类特征

python代码如下：

#!/usr/bin/python
# encoding:utf-8
"""
@author: Bruce
@contact: isunrise@foxmail.com
@file: dtrain.py
@time: 2017/4/3 $ {TIME}
此决策树使用pandas中的dataFrame数据实现，要求dataFrame中指定一列为y标签。调用createtree方法训练。createtree有2个参数
第一个参数中数据DataFrame，第二个参数指定dataFrame中y标签的列名为str类型。
"""

from math import log
import pandas as pd
import operator
import numpy as np
import matplotlib.pyplot as plt

#计算信息熵
def calcshannonent(dataset, re):
    numentries = len(dataset)
    #计算数据集每一分类数量
    l = dataset.columns.tolist()
    k = l.index('Survived')-1
    s_k = pd.pivot_table(dataset, index=re, values=l[k], aggfunc=len)
    #classlabel = set(dataset[re])
    shannonent = 0.0

    #每一分类的信息熵
    for i in list(s_k):
        prob = i / float(numentries)
        shannonent += prob * log(prob, 2)
    return shannonent

#对给定的数据集的指定特征的分类值进行分类
def splitdataset(dataset, axis, value):
    retdataset = dataset[dataset[axis] == value]
    del retdataset[axis]
    return retdataset

#选择最佳分裂特征：
def chooseBestFeatureToSplit(dataset,re):
    #分裂前的信息熵
    baseEntroy = calcshannonent(dataset,re)

    #信息增益及最佳分裂特征初始：
    bestinfogain = 0.0
    bestfeature = dataset.columns[1]

    #对每一特征进行循环
    for f in dataset.columns:
        if f == 'Survived': continue

        #获取当前特征的列表
        featlist = dataset[f]
        #确定有多少分裂值
        uniqueVals = set(featlist)

        #初始化当前特征信息熵为0
        newEntrypoy = 0.0

        #对每一分裂值计算信息熵
        for value in uniqueVals:
            #分裂后的数据集
            subdataset = splitdataset(dataset, f ,value)
            #计算分支的概率
            prob = len(dataset[dataset[f]==value])/float(len(dataset))

            #分裂后信息熵
            newEntrypoy += prob*calcshannonent(subdataset, re)

        # 计算信息增益
        infogain =  newEntrypoy - baseEntroy

        #如果信息增益最大，则替换原信息增益，并记录最佳分裂特征
        if f != 'Survived':
            if (infogain > bestinfogain):
                bestinfogain = infogain
                bestfeature = f

    #返回最佳分裂特征号
    return bestfeature

#确定分枝的主分类
def majority(labellist):
    classcount = {}
    #分类列表中的各分类进行投票
    for vote in labellist:
        if vote not in classcount.keys():
            classcount[vote] =0
        classcount[vote] += 1
    #排序后选择最多票数的分类
    sortedclasscount = sorted(classcount.iteritems(), key=operator.itemgetter(1), reverse=True)
    #返回最多票数的分类的标签
    return sortedclasscount[0][0]

def createtree(dataset, result):
    '有2个参数，第一个是dataFrame类型的数据，第个是字符串类型的y标签'
    #如果数据集的分类只有一个，则返回此分类，无需子树生长
    classlist = list(dataset[result].values)
    if classlist.count(classlist[0]) == len(classlist):
        return classlist[0]

    #如果数据集仅有1列变量，加y标签有2列，则返回此数据集的分类
    if len(dataset.columns) == 2:
        return majority(classlist)

    bestfeat = chooseBestFeatureToSplit(dataset,result)
    mytree = {bestfeat: {}}

    #此节点分裂为哪些分枝
    uniquevals = set(dataset[bestfeat])

    #对每一分枝，递归创建子树
    for value in uniquevals:
        mytree[bestfeat][value] = createtree(splitdataset(dataset, bestfeat,value),result)

    #完成后，返回决策树
    return mytree