代碼
"""
用字典存儲決策樹結構：
{'有自己的房子':{0:{'有工作':{0:'no', 1:'yes'}}, 1:'yes'}}
年齡：0代表青年，1代表中年，2代表老年
有工作：0代表否，1代表是
有自己的房子：0代表否，1代表是
信貸情況：0代表一般，1代表好，2代表非常好
類別（是否給貸款）：no代表否，yes代表是
pickle包可以將決策樹保存下來，方便下次直接調用
"""
from matplotlib.font_manager import FontProperties
import matplotlib.pyplot as plt
from math import log
import operator
import pickle
def createDataSet():
"""
函數說明：創建測試數據集；《統計學習方法》第二版的P71中表5.1的數據
Parameters:
None
Returns:
dataSet - 數據集
labels - 各列特徵的標籤
"""
# 數據集
dataSet = [[0, 0, 0, 0, 'no'],
[0, 0, 0, 1, 'no'],
[0, 1, 0, 1, 'yes'],
[0, 1, 1, 0, 'yes'],
[0, 0, 0, 0, 'no'],
[1, 0, 0, 0, 'no'],
[1, 0, 0, 1, 'no'],
[1, 1, 1, 1, 'yes'],
[1, 0, 1, 2, 'yes'],
[1, 0, 1, 2, 'yes'],
[2, 0, 1, 2, 'yes'],
[2, 0, 1, 1, 'yes'],
[2, 1, 0, 1, 'yes'],
[2, 1, 0, 2, 'yes'],
[2, 0, 0, 0, 'no']]
# 各列特徵的標籤
labels = ['年齡', '有工作', '有自己的房子', '信貸情況']
# 返回數據集及各列特徵的標籤
return dataSet, labels
def calcShannonEnt(dataSet):
"""
函數說明：計算給定數據集的經驗熵（香農熵）
Ent = -SUM（prob*Log2（prob））
Parameters:
dataSet - 數據集
Returns:
shannonEnt - 經驗熵（香農熵）
"""
# 返回數據集的行數
numEntires = len(dataSet)
# 構建標籤（Label）計數的字典
labelCounts = {}
# 對每組特徵向量進行統計
for featVec in dataSet:
# 提取標籤（Label）信息
currentLabel = featVec[-1]
# 如果標籤（Label）沒有放入統計次數的字典，添加進去
if currentLabel not in labelCounts.keys():
# 創建一個新的鍵值對，鍵爲currentLabel值爲0
labelCounts[currentLabel] = 0
# Label計數
labelCounts[currentLabel] += 1
# 經驗熵（香農熵）
shannonEnt = 0.0
# 計算香農熵
for key in labelCounts:
# 選擇該標籤（Label）的概率
prob = float(labelCounts[key]) / numEntires
# 利用公式計算
shannonEnt -= prob*log(prob, 2)
# 返回經驗熵（香農熵）
return shannonEnt
def splitDataSet(dataSet, axis, value):
"""
函數說明：按照指定特徵的值劃分數據集
Parameters:
dataSet - 待劃分的數據集
axis - 指定劃分數據集的特徵
values - 指定特徵的值劃
Returns:
retDataSet - 按照
"""
# 創建返回的數據集列表
retDataSet = []
# 遍歷數據集的每一行
for featVec in dataSet:
if featVec[axis] == value:
# 去掉axis特徵列
reducedFeatVec = featVec[:axis]
# extend() 函數用於在列表末尾一次性追加另一個序列中的多個值（用新列表擴展原來的列表）。
reducedFeatVec.extend(featVec[axis+1:])
# 將劃分後的數據添加到數據集
retDataSet.append(reducedFeatVec)
# 返回劃分後的數據集
return retDataSet
def chooseBestFeatureToSplit(dataSet):
"""
核心算法：
ID3：選取信息增益最大的特徵（此函數用的算法）
C4.5：選取信息增益比最大的特徵
函數說明：選擇最優特徵
Gain(D,g) = Ent(D) - SUM(|Dv|/|D|)*Ent(Dv)
Parameters:
dataSet - 數據集
Returns:
bestFeature - 信息增益最大的（最優）特徵的索引值
"""
# 特徵數量
numFeatures = len(dataSet[0]) - 1
# 計算數據集的香農熵
baseEntropy = calcShannonEnt(dataSet)
# 信息增益
bestInfoGain = 0.0
# 最優特徵的索引值
bestFeature = -1
# 遍歷數據集中各列特徵
for i in range(numFeatures):
# 獲取dataSet中第i個所有特徵值，存放在featList列表中（列表生成式）
featList = [example[i] for example in dataSet]
# 創建set集合{}，元素不可重複，重複的元素均被刪掉
uniqueVals = set(featList)
# 經驗條件熵
newEntropy = 0.0
# 計算信息增益
for value in uniqueVals:
# subDataSet爲按照指定特徵列給定的特徵值進行劃分後的子集
subDataSet = splitDataSet(dataSet, i, value)
# 計算子集的概率
prob = len(subDataSet) / float(len(dataSet))
# 根據公式計算經驗條件熵
newEntropy += prob * calcShannonEnt(subDataSet)
# 信息增益
infoGain = baseEntropy - newEntropy
# 打印每個特徵的信息增益
print("第%d個特徵的增益爲%.3f" % (i, infoGain))
# 計算信息增益
if(infoGain > bestInfoGain):
# 更新信息增益，找到最大的信息增益
bestInfoGain = infoGain
# 記錄信息增益最大的特徵的索引值
bestFeature = i
# 返回信息增益最大的特徵的索引值
return bestFeature
def majorityCnt(classList):
"""
函數說明：統計classList中出現次數最多的元素（類標籤）
服務於遞歸的第2個終止條件
Parameters:
classList - 類標籤列表
Returns:
sortedClassCount[0][0] - 出現次數最多的元素（類標籤）
"""
classCount = {}
# 統計classList中每個元素出現的次數
for vote in classList:
if vote not in classCount.keys():
classCount[vote] = 0
classCount[vote] += 1
# 根據字典的值降序排序
# operator.itemgetter(1)獲取對象的第1列的值
sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
# 返回classList中出現次數最多的元素
return sortedClassCount[0][0]
def createTree(dataSet, labels, featLabels):
"""
函數說明：創建決策樹（ID3算法）
遞歸有兩個終止條件：1、所有的類標籤完全相同，直接返回類標籤
2、用完所有標籤但是得不到唯一類別的分組，即特徵不夠用，挑選出現數量最多的類別作爲返回
Parameters:
dataSet - 訓練數據集
labels - 分類屬性標籤
featLabels - 存儲選擇的最優特徵標籤
Returns:
myTree - 決策樹
"""
# 取分類標籤（是否放貸：yes or no）
classList = [example[-1] for example in dataSet]
# 遞歸終止條件1：如果類別完全相同則停止繼續劃分
if classList.count(classList[0]) == len(classList):
return classList[0]
# 遞歸終止條件2：遍歷完所有特徵時返回出現次數最多的類標籤
if len(dataSet[0]) == 1:
return majorityCnt(classList)
# 選擇最優特徵,
bestFeat = chooseBestFeatureToSplit(dataSet)
# 最優特徵的標籤
bestFeatLabel = labels[bestFeat]
featLabels.append(bestFeatLabel)
# 根據最優特徵的標籤生成樹
myTree = {bestFeatLabel: {}}
# 刪除已經使用的特徵標籤
del(labels[bestFeat])
# 得到訓練集中所有最優解特徵的屬性值
featValues = [example[bestFeat] for example in dataSet]
# 去掉重複的屬性值
uniqueVals = set(featValues)
# 遍歷特徵，創建決策樹
for value in uniqueVals:
myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), labels, featLabels)
return myTree
def getNumLeafs(myTree):
"""
函數說明：獲取決策樹葉子結點的數目
Parameters:
myTree - 決策樹
Returns:
numLeafs - 決策樹的葉子結點的數目
"""
# 初始化葉子
numLeafs = 0
# python3中myTree.keys()返回的是dict_keys,不是list,所以不能用
# myTree.keys()[0]的方法獲取結點屬性，可以使用list(myTree.keys())[0]
# next() 返回迭代器的下一個項目 next(iterator[, default])
firstStr = next(iter(myTree))
# 獲取下一組字典
secondDict = myTree[firstStr]
for key in secondDict.keys():
# 測試該結點是否爲字典，如果不是字典，代表此節點爲葉子結點
if type(secondDict[key]).__name__ == 'dict':
numLeafs += getNumLeafs(secondDict[key])
else:
numLeafs += 1
return numLeafs
def getTreeDepth(myTree):
"""
函數說明：獲取決策樹的層數
Parameters:
myTree - 決策樹
Returns:
maxDepth - 決策樹的層數
"""
# 初始化決策樹深度
maxDepth = 0
# python3中myTree.keys()返回的是dict_keys,不是list,所以不能用
# myTree.keys()[0]的方法獲取結點屬性，可以使用list(myTree.keys())[0]
# next() 返回迭代器的下一個項目 next(iterator[, default])
firstStr = next(iter(myTree))
# 獲取下一個字典
secondDict = myTree[firstStr]
for key in secondDict.keys():
# 測試該結點是否爲字典，如果不是字典，代表此節點爲葉子結點
if type(secondDict[key]).__name__ == 'dict':
thisDepth = 1 + getTreeDepth(secondDict[key])
else:
thisDepth = 1
# 更新最深層數
if thisDepth > maxDepth:
maxDepth = thisDepth
# 返回決策樹的層數
return maxDepth
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
"""
函數說明：繪製結點
Parameters:
nodeTxt - 結點名
centerPt - 文本位置
parentPt - 標註的箭頭位置
nodeType - 結點格式
Returns:
None
"""
# 定義箭頭格式
arrow_args = dict(arrowstyle="<-")
# 設置中文字體
font = FontProperties(fname=r"C:\Windows\Fonts\simsun.ttc", size=14)
# 繪製結點createPlot.ax1創建繪圖區
# annotate是關於一個數據點的文本
# nodeTxt爲要顯示的文本，centerPt爲文本的中心點，箭頭所在的點，parentPt爲指向文本的點
createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
xytext=centerPt, textcoords='axes fraction',
va='center', ha='center', bbox=nodeType,
arrowprops=arrow_args, FontProperties=font)
def plotMidText(cntrPt, parentPt, txtString):
"""
函數說明：標註有向邊屬性值
Parameters:
cntrPt、parentPt - 用於計算標註位置
txtString - 標註內容
Returns:
None
"""
# 計算標註位置（箭頭起始位置的中點處）
xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)
def plotTree(myTree, parentPt, nodeTxt):
"""
函數說明：繪製決策樹
Parameters:
myTree - 決策樹（字典）
parentPt - 標註的內容
nodeTxt - 結點名
Returns:
None
"""
# 設置結點格式boxstyle爲文本框的類型，sawtooth是鋸齒形，fc是邊框線粗細
decisionNode = dict(boxstyle="sawtooth", fc="0.8")
# 設置葉結點格式
leafNode = dict(boxstyle="round4", fc="0.8")
# 獲取決策樹葉結點數目，決定了樹的寬度
numLeafs = getNumLeafs(myTree)
# 獲取決策樹層數
depth = getTreeDepth(myTree)
# 下個字典
firstStr = next(iter(myTree))
# 中心位置
cntrPt = (plotTree.xoff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yoff)
# 標註有向邊屬性值
plotMidText(cntrPt, parentPt, nodeTxt)
# 繪製結點
plotNode(firstStr, cntrPt, parentPt, decisionNode)
# 下一個字典，也就是繼續繪製結點
secondDict = myTree[firstStr]
# y偏移
plotTree.yoff = plotTree.yoff - 1.0 / plotTree.totalD
for key in secondDict.keys():
# 測試該結點是否爲字典，如果不是字典，代表此結點爲葉子結點
if type(secondDict[key]).__name__ == 'dict':
# 不是葉結點，遞歸調用繼續繪製
plotTree(secondDict[key], cntrPt, str(key))
# 如果是葉結點，繪製葉結點，並標註有向邊屬性值
else:
plotTree.xoff = plotTree.xoff + 1.0 / plotTree.totalW
plotNode(secondDict[key], (plotTree.xoff, plotTree.yoff), cntrPt, leafNode)
plotMidText((plotTree.xoff, plotTree.yoff), cntrPt, str(key))
plotTree.yoff = plotTree.yoff + 1.0 / plotTree.totalD
def createPlot(inTree):
"""
函數說明：創建繪圖面板
Parameters:
inTree - 決策樹（字典）
Returns:
None
"""
# 創建fig
fig = plt.figure(1, facecolor="white")
# 清空fig
fig.clf()
axprops = dict(xticks=[], yticks=[])
# 去掉x、y軸
createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
# 獲取決策樹葉結點數目
plotTree.totalW = float(getNumLeafs(inTree))
# 獲取決策樹層數
plotTree.totalD = float(getTreeDepth(inTree))
# x偏移
plotTree.xoff = -0.5 / plotTree.totalW
plotTree.yoff = 1.0
# 繪製決策樹
plotTree(inTree, (0.5, 1.0), '')
# 顯示繪製結果
plt.show()
def classify(inputTree, featLabels, testVec):
"""
函數說明：使用決策樹分類
Parameters:
inputTree - 已經生成的決策樹
featLabels - 存儲選擇的最優特徵標籤
testVec - 測試數據列表，順序對應最優特徵標籤
Returns:
classLabel - 分類結果
"""
# 獲取決策樹結點
firstStr = next(iter(inputTree))
# 下一個字典
secondDict = inputTree[firstStr]
featIndex = featLabels.index(firstStr)
for key in secondDict.keys():
if testVec[featIndex] == key:
if type(secondDict[key]).__name__ == 'dict':
classLabel = classify(secondDict[key], featLabels, testVec)
else:
classLabel = secondDict[key]
return classLabel
def storeTree(inputTree, filename):
"""
函數說明：存儲決策樹
Parameters:
inputTree - 已經生成的決策樹
filename - 決策樹的存儲文件名
Returns:
None
"""
with open(filename, 'wb') as fw:
pickle.dump(inputTree, fw)
def grabTree(filename):
"""
函數說明：讀取決策樹
Parameters:
filename - 決策樹的存儲文件名
Returns:
pickle.load(fr) - 決策樹字典
"""
fr = open(filename, 'rb')
return pickle.load(fr)
def main():
# 生成數據集及各列特徵的標籤
dataSet, labels = createDataSet()
# 創建featLabels：決策樹最佳特徵劃分先後順序的列表
featLabels = []
myTree = createTree(dataSet, labels, featLabels)
print("決策樹最佳特徵劃分先後順序的", featLabels)
# storeTree(myTree, 'classifierStorage.txt')
# myTree = grabTree('classifierStorage.txt')
# print(myTree)
# 測試數據
testVec = [0, 1, 1, 1]
result = classify(myTree, featLabels, testVec)
if result == 'yes':
print('放貸')
if result == 'no':
print('不放貸')
print(myTree)
createPlot(myTree)
# print(dataSet)
# print(calcShannonEnt(dataSet))
print("最優特徵索引值:" + str(chooseBestFeatureToSplit(dataSet)))
if __name__ == '__main__':
main()
運行結果
第0個特徵的增益爲0.083
第1個特徵的增益爲0.324
第2個特徵的增益爲0.420
第3個特徵的增益爲0.363
第0個特徵的增益爲0.252
第1個特徵的增益爲0.918
第2個特徵的增益爲0.474
決策樹最佳特徵劃分先後順序的 ['有自己的房子', '有工作']
放貸
{'有自己的房子': {0: {'有工作': {0: 'no', 1: 'yes'}}, 1: 'yes'}}
第0個特徵的增益爲0.083
第1個特徵的增益爲0.324
第2個特徵的增益爲0.420
第3個特徵的增益爲0.363
最優特徵索引值:2
