目录

一、前言

二、代码实现逻辑

        构造树模块

        选择最好特征模块

        计算熵值模块

        切分数据集模块

        当前样本中多数类别模块

三、可视化拓展

 四、结果展示+完整代码

一、前言

        本文需要读者有对决策树有一定的基础，可以参考决策树原理(决策树算法概述，熵，信息增益，信息增益率，gini系数，剪枝，回归、分类任务解决)

二、代码实现逻辑

        构造树模块

        （学过数据结构的都知道，构造树最好的方法是递归）

        1.判断是否需要建树：如果当前节点所有样本的标签相同，不需要建树，如果所有特征都用完还是没有完全分类好，则分类结果采取需要少数服从多数的策略。

        2.把最好的那个特征选出来用来当作根节点

        3.根据根节点的不同特征值进行分叉

        4.在数据集中把以根节点为特征的特征值去掉（更新数据集）

        5.在特征值里循环递归建树

        6.返回树

        注意：采用字典嵌套的形式来存储树，featLabels表示根节点的值，可以根据先后顺序把特征值存储起来。

        

def crecateTree(dataset,labels,featLabels):
	'''

	:param dataset: 数据集
	:param labels: 判断当前节点是否需要再分
	:param featLabels: 根节点的值
	:return:
	'''
	classList = [example[-1] for example in dataset] #当前节点的所有样本的标签
	if classList.count(classList[0]) == len(classList): #判断所有标签是否一致
		return classList[0]
	if len(dataset[0]) == 1: #只剩下一列特征值
		return majorityCnt(classList) #返回主要特征
	bestFeature = chooseBestFeatureToSplit(dataset) #得到最好特征的索引
	bestFeatureLabel = labels[bestFeature]
	featLabels.append(bestFeatureLabel)
	myTree = {bestFeatureLabel:{}} #用字典来存储树，嵌套
	del labels[bestFeature] #删除特征值
	featValue = [example[bestFeature] for example in dataset] #得到根节点特征值
	uniqueVals = set(featValue)# 有几个不同的特征值，树分几个叉
	for value in uniqueVals: #递归调用
		sublabels = labels[:]
		myTree[bestFeatureLabel][value] = crecateTree(splitDataSet(dataset,bestFeature,value),sublabels,featLabels)
	return myTree

        选择最好特征模块

        需要把每个特征都遍历一遍，选择信息增益最大的那个特征

        

def chooseBestFeatureToSplit(dataset): #核心，熵值计算
	numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去label
	baseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值
	bestInfoGain = 0 #最好的信息增益
	bestFeature = -1 #最好的特征
	for i in range(numFeatures):
		featList = [example[i] for example in dataset] #当前的特征列
		uniqueVals = set(featList) #特征值的类别
		newEntropy = 0
		for val in uniqueVals:
			subDataSet = splitDataSet(dataset,i,val)
			prob = len (subDataSet) / float(len(dataset))
			newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值
		infoGain = baseEntropy - newEntropy
		if(infoGain > bestInfoGain):
			bestInfoGain = infoGain
			bestFeature = i
	return bestFeature

        计算熵值模块

        把需要的概率算出来

def calcShannonEnt(dataset):#熵值计算
	numexamples = len(dataset)
	labelCount = {}
	for featVec in dataset:
		currentlabel = featVec[-1]
		if currentlabel not in labelCount.keys():
			labelCount[currentlabel] = 0
		labelCount[currentlabel] += 1

	shannonEnt = 0
	for key in labelCount:
		prop = float(labelCount[key]/numexamples) #概率值
		shannonEnt -= prop*log(prop,2)  #熵值
	return shannonEnt

        切分数据集模块

        每次进行划分后都需要数据切分，包括去掉根节点特征的那一列

def splitDataSet(dataset,axis,val): #切分数据集，把根节点的那一特征列去掉
	retDataSet = []
	for featVec in dataset:
		if featVec[axis] == val:
			reducedFeatVec = featVec[:axis]
			reducedFeatVec.extend(featVec[axis+1:]) #用切片和拼接把第axis列切掉
			retDataSet.append(reducedFeatVec)
	return retDataSet

        当前样本中多数类别模块

        当所有的特征都用完后还不能完全划分，采取少数服从多数策略

def majorityCnt(classList):		#当前多数类别是哪一个
	classCount = {}
	for vote in classList:
		if vote not in classCount.keys():
			classCount[vote] = 0
		classCount[vote] += 1
	sortedClassCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True) #排序
	return sortedClassCount[0][0]

三、可视化拓展

        这个不是重点，重要的是掌握递归建树的思想

        

def getNumLeafs(myTree):
	numLeafs = 0
	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):
	maxDepth = 0
	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):
	arrow_args = dict(arrowstyle="<-")
	font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)

	createPlot.axl.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):
	xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
	yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
	createPlot.axl.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def plotTree(myTree, parentPt, nodeTxt):
	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]
	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):
	fig = plt.figure(1, facecolor='white')													#创建fig
	fig.clf()																				#清空fig
	axprops = dict(xticks=[], yticks=[])
	createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    							#去掉x、y轴
	plotTree.totalW = float(getNumLeafs(inTree))											#获取决策树叶结点数目
	plotTree.totalD = float(getTreeDepth(inTree))											#获取决策树层数
	plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0							#x偏移
	plotTree(inTree, (0.5,1.0), '')															#绘制决策树
	plt.show()

 四、结果展示+完整代码

# -*- coding: UTF-8 -*-
from matplotlib.font_manager import FontProperties
import matplotlib.pyplot as plt
from math import log
import operator



def createDataSet():
	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 = ['F1-AGE', 'F2-WORK', 'F3-HOME', 'F4-LOAN']		
	return dataSet, labels

def crecateTree(dataset,labels,featLabels):
	'''

	:param dataset: 数据集
	:param labels: 判断当前节点是否需要再分
	:param featLabels: 节点的值
	:return:
	'''
	classList = [example[-1] for example in dataset] #当前节点的所有样本的标签
	if classList.count(classList[0]) == len(classList): #判断所有标签是否一致
		return classList[0]
	if len(dataset[0]) == 1: #只剩下一列特征值
		return majorityCnt(classList) #返回主要特征
	bestFeature = chooseBestFeatureToSplit(dataset) #得到最好特征的索引
	bestFeatureLabel = labels[bestFeature]
	featLabels.append(bestFeatureLabel)
	myTree = {bestFeatureLabel:{}} #用字典来存储树，嵌套
	del labels[bestFeature] #删除特征值
	featValue = [example[bestFeature] for example in dataset] #得到根节点特征值
	uniqueVals = set(featValue)# 有几个不同的特征值，树分几个叉
	for value in uniqueVals: #递归调用
		sublabels = labels[:]
		myTree[bestFeatureLabel][value] = crecateTree(splitDataSet(dataset,bestFeature,value),sublabels,featLabels)
	return myTree

def majorityCnt(classList):		#当前多数类别是哪一个
	classCount = {}
	for vote in classList:
		if vote not in classCount.keys():
			classCount[vote] = 0
		classCount[vote] += 1
	sortedClassCount = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True) #排序
	return sortedClassCount[0][0]

def chooseBestFeatureToSplit(dataset): #核心，熵值计算
	numFeatures = len(dataset[0]) - 1 #特征的个数 随便一列减去label
	baseEntropy = calcShannonEnt(dataset) #计算当前什么都不做的熵值
	bestInfoGain = 0 #最好的信息增益
	bestFeature = -1 #最好的特征
	for i in range(numFeatures):
		featList = [example[i] for example in dataset] #当前的特征列
		uniqueVals = set(featList) #特征值的类别
		newEntropy = 0
		for val in uniqueVals:
			subDataSet = splitDataSet(dataset,i,val)
			prob = len (subDataSet) / float(len(dataset))
			newEntropy += prob * calcShannonEnt(subDataSet) # 选择特征后的熵值
		infoGain = baseEntropy - newEntropy
		if(infoGain > bestInfoGain):
			bestInfoGain = infoGain
			bestFeature = i
	return bestFeature


def splitDataSet(dataset,axis,val): #切分数据集，把根节点的那一特征列去掉
	retDataSet = []
	for featVec in dataset:
		if featVec[axis] == val:
			reducedFeatVec = featVec[:axis]
			reducedFeatVec.extend(featVec[axis+1:]) #用切片和拼接把第axis列切掉
			retDataSet.append(reducedFeatVec)
	return retDataSet

def calcShannonEnt(dataset):#熵值计算
	numexamples = len(dataset)
	labelCount = {}
	for featVec in dataset:
		currentlabel = featVec[-1]
		if currentlabel not in labelCount.keys():
			labelCount[currentlabel] = 0
		labelCount[currentlabel] += 1

	shannonEnt = 0
	for key in labelCount:
		prop = float(labelCount[key]/numexamples) #概率值
		shannonEnt -= prop*log(prop,2)  #熵值
	return shannonEnt


def getNumLeafs(myTree):
	numLeafs = 0
	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):
	maxDepth = 0
	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):
	arrow_args = dict(arrowstyle="<-")
	font = FontProperties(fname=r"c:\windows\fonts\simsunb.ttf", size=14)

	createPlot.axl.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):
	xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
	yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
	createPlot.axl.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def plotTree(myTree, parentPt, nodeTxt):
	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]
	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):
	fig = plt.figure(1, facecolor='white')													#创建fig
	fig.clf()																				#清空fig
	axprops = dict(xticks=[], yticks=[])
	createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)    							#去掉x、y轴
	plotTree.totalW = float(getNumLeafs(inTree))											#获取决策树叶结点数目
	plotTree.totalD = float(getTreeDepth(inTree))											#获取决策树层数
	plotTree.xOff = -0.5/plotTree.totalW; plotTree.yOff = 1.0							#x偏移
	plotTree(inTree, (0.5,1.0), '')															#绘制决策树
	plt.show()

if __name__ == '__main__':
	dataSet,labels = createDataSet()
	featLabels = []
	myTree = crecateTree(dataSet, labels, featLabels)
	print(featLabels)
	createPlot(myTree)

 

        选择两个特征建树

        可视化结果：

# 决策树 # 算法 # 机器学习 # python
Logo
2048 AI社区

有“AI”的1024 = 2048，欢迎大家加入2048 AI社区

更多推荐

cover

AI 绘画的未来趋势与发展前景

avatar 2048 AI社区
cover

CVPRW 2025 | Qualcomm AI提出基于不确定性的光流与立体深度估计改进方法！

avatar 2048 AI社区

剖析：量子计算如何优化提示工程架构流程

当我们用ChatGPT写邮件、用Claude分析论文、用Gemini生成代码时，**提示工程（Prompt Engineering）**早已成为连接人类意图与AI能力的“翻译器”。长文档问答时，要么截断上下文丢失关键信息，要么拆分文档破坏逻辑；调试prompt参数（温度、top-k、指令模板）时，试错成本高到“调一天prompt，跑一次模型”；多轮对话中，AI经常“健忘”——前面提到的细节后面全丢

avatar 2048 AI社区