# 生成决策树
from math import log
import operator
import pickle
def createDataSet():
dataSet = [[1, 1, 'yes'],
[1, 1, 'yes'],
[1, 0, 'no'],
[0, 1, 'no'],
[0, 1, 'no']]
labels = ['no surfacing', 'flippers']
return dataSet, labels
def calcShannonEnt(data_set):
num_entries = len(data_set) # 行数
label_counts = {}
for feat_vec in data_set: # 循环每行
current_label = feat_vec[-1] # 矩阵最后一列为标签
if current_label not in label_counts.keys(): # 若字典内没标签
label_counts[current_label] = 0 # 将标签：0，以key：value 形式存进字典中
label_counts[current_label] += 1 # 字典内有该标签，则该标签count+1
shannon_ent = 0.0
for key in label_counts:
prob = float(label_counts[key]) / num_entries # 求出每个标签的概率
shannon_ent -= prob * log(prob, 2) # 熵 = 所有元素的-(p(x1)log(x1,2))求和
return shannon_ent # 返回熵
def splitDataSet(data_set, axis, value):
'''
涮选出根据axis与value划分后的列表
:param data_set: 数据集
:param axis: 根据axis去除列
:param value: 根据value去除行
:return:
'''
ret_data_set = [] # 创建新列表
for feat_vec in data_set: # 循环每行
if feat_vec[axis] == value:
reduced_feat_vec = feat_vec[:axis]
reduced_feat_vec.extend(feat_vec[axis + 1:]) # 这两步等于删除feat_vec[axis]，但是为了不影响原始数据，所以分两步切片合并
ret_data_set.append(reduced_feat_vec)
return ret_data_set # 返回划分后的数据集
def chooseBestFeatureToSplit(data_set): # 找到最大熵对应的索引
num_features = len(data_set[0]) - 1 # data_set最后一个元素已用
base_entropy = calcShannonEnt(data_set) # 整个数据集的熵，保存最初的无序度量值
best_info_gain = 0.0
best_feature = -1
for i in range(num_features):
feat_list = [example[i] for example in data_set] # 利用列表生成式产生新列表
unique_vals = set(feat_list) # 利用集合去重
new_entropy = 0.0
for value in unique_vals: # 遍历去重后的唯一属性值
sub_data_set = splitDataSet(data_set, i, value) # 已唯一属性值为value划分出数据集
prob = len(sub_data_set) / float(len(data_set)) # 计算该数据集的概率
new_entropy += prob * calcShannonEnt(sub_data_set) # 计算熵
info_gain = base_entropy - new_entropy
if info_gain > best_info_gain: # 判断是否大于最佳熵
best_info_gain = info_gain
best_feature = i
return best_feature # 返回最大熵对应的索引
def majorityCnt(class_list):
class_count = {}
for vote in class_list:
if vote not in class_count.keys():
class_count[vote] = 0
class_count[vote] += 1
sorted_class_count = sorted(class_count.items(), key=operator.itemgetter(1), reverse=True)
return sorted_class_count[0][0] # 返回次数最多的标签
def createTree(data_set, labels):
class_list = [example[-1] for example in data_set] # 拿到各个标签的列表
if class_list.count(class_list[0]) == len(class_list): # 所有类标签完全相同时
return class_list[0] # 返回类标签,return是结束本次递归循环并返回值
if len(data_set[0]) == 1: # 遍历玩所有标签，仍然没有把数据集划分为唯一类别
return majorityCnt(class_list) # 返回次数最多的标签
best_feat = chooseBestFeatureToSplit(data_set) # 找到最大熵对应的索引
best_feat_label = labels[best_feat] # 最合适的标签
my_tree = {best_feat_label: {}} # 创建tree字典
del (labels[best_feat])
feat_values = [example[best_feat] for example in data_set] # 得到最大熵对应列的所有属性值
unique_vals = set(feat_values) # 去除相同属性值
for value in unique_vals:
sub_labels = labels[:] # 保证for期间
my_tree[best_feat_label][value] = createTree(splitDataSet(data_set, best_feat, value),
sub_labels) # 给字典增加键为value，值为返回值的新键值
return my_tree
def classify(input_tree, feat_tabels, test_vec):
first_str = list(input_tree.keys())[0] # 决策树顶端key
second_dict = input_tree[first_str] # 第二层字典
feat_index = feat_tabels.index(first_str) # 最佳分类对应的索引
for key in second_dict.keys():
if test_vec[feat_index] == key:
if type(second_dict[key]).__name__ == 'dict':
class_label = classify(second_dict[key], feat_tabels, test_vec)
else:
class_label = second_dict[key]
return class_label
def storeTree(input_tree, filename):
with open(filename, 'wb') as fw:
pickle.dump(input_tree, fw)
fw.close()
def grabTree(filename):
with open(filename, 'rb') as fr:
return pickle.load(fr)
绘制决策树
import matplotlib.pyplot as plt
# 把决策树的字典放入createPlot(dict) 即可
decision_node = dict(boxstyle='sawtooth', fc='0.8') # 填文本底色
leaf_node = dict(boxstyle='round4', fc='0.8')
arrow_args = dict(arrowstyle='<-')
def plotNode(node_text, center_pt, parent_pt, node_type): # 绘制带箭头的注解
createPlot.ax1.annotate(node_text, xy=parent_pt, xycoords='axes fraction',
xytext=center_pt, textcoords='axes fraction', va='center',
ha='center', bbox=node_type, arrowprops=arrow_args)
def getNumLeafs(my_tree): # 找出决策树中没分叉的数量
num_leafs = 0
first_str = list(my_tree.keys())[0] # 取出决策树字典中的最佳值
second_dict = my_tree[first_str] # 取出第二个字典
for key in second_dict.keys(): # 遍历字典的key
if type(second_dict[key]).__name__ == 'dict': # 如果key对应的value是字典，则递归继续分
num_leafs += getNumLeafs(second_dict[key])
else:
num_leafs += 1 # 不是则+1
return num_leafs
def getTreeDepth(my_tree): # 找出决策树的层数，第一层不算
max_depth = 0
first_str = list(my_tree.keys())[0] # 取出决策树字典中的最佳值
second_dict = my_tree[first_str] # 取出第二个字典
for key in second_dict.keys(): # 遍历字典的key
if type(second_dict[key]).__name__ == 'dict':
this_depth = 1 + getTreeDepth(second_dict[key])
else:
this_depth = 1
if this_depth > max_depth:
max_depth = this_depth
return max_depth
def retrieveTree(i):
list_of_trees = [{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},
{'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yse'}}, 1: 'no'}}}}]
return list_of_trees[i]
def plotMidText(cntr_pt, parent_pt, txt_string): # 在两个节点之间的线上写字
xMid = (parent_pt[0] - cntr_pt[0]) / 2.0 + cntr_pt[0]
yMid = (parent_pt[1] - cntr_pt[1]) / 2 + cntr_pt[1]
createPlot.ax1.text(xMid, yMid, txt_string)
def createPlot(in_tree):
fig = plt.figure(1, facecolor='white')
fig.clf() # 清空绘图区
axprops = dict(xticks=[], yticks=[])
createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
global plotTree_w
global plotTree_d
plotTree_w = float(getNumLeafs(in_tree))
plotTree_d = float(getTreeDepth(in_tree))
plotTree.xOff = -0.5 / plotTree_w
plotTree.yOff = 1.0
plotTree(in_tree, (0.5, 1.0), '')
plt.show()
def plotTree(my_tree, parent_pt, node_txt): # 画树
num_leafs = getNumLeafs(my_tree)
depth = getTreeDepth(my_tree)
first_str = list(my_tree.keys())[0]
cntr_pt = (plotTree.xOff + (1.0 + float(num_leafs)) / 2.0 / plotTree_w, plotTree.yOff)
plotMidText(cntr_pt, parent_pt, node_txt)
plotNode(first_str, cntr_pt, parent_pt, decision_node)
second_dict = my_tree[first_str]
plotTree.yOff = plotTree.yOff - 1.0 / plotTree_d # 减少y的值，将树的总深度平分，每次减少移动一点(向下，因为树是自顶向下画的）
for key in second_dict.keys():
if type(second_dict[key]).__name__ == 'dict':
plotTree(second_dict[key], cntr_pt, str(key))
else:
plotTree.xOff = plotTree.xOff + 1.0 / plotTree_w
plotNode(second_dict[key], (plotTree.xOff, plotTree.yOff), cntr_pt, leaf_node)
plotMidText((plotTree.xOff, plotTree.yOff), cntr_pt, str(key))
plotTree.yOff = plotTree.yOff + 1.0 / plotTree_d