上一篇文章机器学习算法复习手册——决策树在复习完基本概念之后，我给自己挖了一个坑：用python写一个决策树出来（注意，不是sklearn调包）。虽然说这个东西在几年前我写过一次，但又写一次，发现很多地方还有挺折磨我的。今天，就来填这个坑，分享一下我写的很垃圾的ID3决策树算法。
注：
- 这里展示的是ID3算法的大致思想
- 没有写树的剪枝算法，只有“特征选择”和“树的生成”
- 没有经过细致的优化，没有对照网上的其他写法
一、数据集
这里为了方便，我搞了一份十分简陋的“西瓜分类数据集”：
总共就17个样本，配得上“简陋”二字。 所以，我们的实验中就不涉及训练集测试集划分了，只要能根据这个数据集，生成一个符合它的规则的决策树，就算是成功了。 数据集文件可以通过公众号后台回复【简陋西瓜数据集】获取。
二、特征选择算法
ID3算法，使用的是“信息增益”的标准来选择特征，这个主要分两步：
- 熵的计算
- 条件熵的计算
算法可以回顾一下上一篇文章，这里不赘述，直接上代码：
代码语言：javascript复制import math
# 熵的计算
def entropy(p_list):
E = 0
for p in p_list:
if p == 0:
E = 0
else:
E = -p*math.log2(p)
return E
def cal_info_gain(dataset,feature):
# 先计算dataset的熵：
N = len(dataset)
n1 = len(dataset[dataset["好"]=="是"])
E = entropy([n1/N,(N-n1)/N])
# 再计算条件熵：
# 需要知道：feature的每一个值在dataset中的个数，以及其中各类别的个数
feature_count = {}
for fv,c in zip(list(dataset[feature]),list(dataset["好"])):
if fv not in feature_count:
feature_count[fv] = [0,0] # 创建一个list，存放各类别的个数
if c == "是":
feature_count[fv][1] = 1
else:
feature_count[fv][0] = 1
conditional_E = 0
for v in feature_count:
N_v = sum(feature_count[v]) # N_v：当前feature value的个数
p_list = []
for N_vc in feature_count[v]: # N_vc：feature value中各个类别的个数
p_list.append(N_vc/N_v)
conditional_E = N_v/N * entropy(p_list)
return E - conditional_E
上面的注释部分应该写的比较清楚了，反正主要思想就是“数数”，把数字数清楚了，就都好算了。
具体思想就是：
给定了一个feature（比如“触感”），我希望统计出这个feature的各个值（“硬滑”和“软粘”）各自占数据集的多少，以及每一个值中，数据各个类别的多少。
上面代码里面有一个关键性的数据统计字典：feature_count，我们对feature=”触感”进行统计之后，得到的结果是：
feature_count = {'硬滑': [6, 6], '软粘': [3, 2]}
意思就是在“硬滑”这个特征值下，有6个坏瓜6个好瓜，“软粘”这个值下，有3个坏瓜2个好瓜。那么条件熵就用公式一套即可。
三、树的生成
很明显，我们需要用一个递归函数来生成一棵树。 一个重要的问题是：用什么数据结构存储这个树呢？ 思来想去，发现字典这种非线性结构比较好。线性的数据结构是无法存储树这种东西的。
思路是这样的：
- 一个选择了某feature来划分，该feature就作为字典的一个key；
- 一个feature会有多个feature values，所以这个feature后面还得接一个字典，每一个feature value都是一个key；
- 一个feature value下面，可能唯一确定了一个叶节点了，所以可以接一个标签值label，终止； 也有可能还要继续划分，所以就要确定新的划分feature，这个新的feature就作为key，继续回到步骤1。
所以想象中，我们希望得到这样的一个层次字典：
代码语言：javascript复制{'纹理': {'模糊': '否',
'清晰': {'根蒂': {'硬挺': '否',
'稍蜷': {'色泽': {'乌黑': '否', '青绿': '是'}},
'蜷缩': '是'}},
'稍糊': {'触感': {'硬滑': '否', '软粘': '是'}}}}
当然，我们事先不必画出这么一个字典，太费劲了（这个字典是我程序写完之后生成的）。我们只用脑补一下这个结构就行了。
有了这样的数据结构，程序就不难写了，上代码：
代码语言：javascript复制def tree_generation(dataset,features):
# 终止条件：
if len(set(dataset["好"])) == 1: # 数据集里所有样本都属于同一类
return list(dataset["好"])[0] # 返回类别
if len(features) == 0: # 没有剩余的特征了
if len(dataset[dataset["好"]=="是"])>len(dataset[dataset["好"]=="否"]):
return "是"
else:
return "否"
# 其他情况：
info_gains = []
for feature in features:
info_gains.append(cal_info_gain(dataset,feature))
best_index = info_gains.index(max(info_gains))
best_feature = features[best_index]
features.remove(best_feature)
remaining_features = features[:] # 不能使用features，要复制一份
if max(info_gains)<=0:# 说明所有的特征都没有区分度
if len(dataset[dataset["好"]=="是"])>len(dataset[dataset["好"]=="否"]):
return "是"
else:
return "否"
else:
my_tree = {best_feature:{}}
feature_values = set(dataset[best_feature])
for fv in feature_values:
sub_dataset = dataset[dataset[best_feature]==fv]
my_tree[best_feature][fv] = tree_generation(sub_dataset,remaining_features)
return my_tree
上面的程序，先讨论了三种特殊情况：
- 样本都同一类了
- 没有剩下的特征了
- 有特征，但没有区分度了 这些都是直接返回label，是递归程序的终止条件。
然后，就是正常的特征选择、划分的流程了：
- 每个特征算个info gain，大家一起碰一碰
- 最好的特征，每个特征值对数据集进行分割，产生子数据集，丢入递归程序中
验证一下：
代码语言：javascript复制feature_names = ['色泽','根蒂','敲声','纹理','脐部','触感']
t = tree_generation(df,feature_names)
pprint(t)
得到的输出就是上面给的字典：
代码语言：javascript复制{'纹理': {'模糊': '否',
'清晰': {'根蒂': {'硬挺': '否',
'稍蜷': {'色泽': {'乌黑': '否', '青绿': '是'}},
'蜷缩': '是'}},
'稍糊': {'触感': {'硬滑': '否', '软粘': '是'}}}}
就这样，我们决策树最核心的部分就写好了，数了数，就27行代码，如何把if else语句写成一行式就更短了。
四、如何画出来？
字典形式总归是难看的，我们希望能用一个生动形象的真正的树给它描绘一下。 这里尝试使用matplotlib这个我觉得很难用，每次用都要到处查文档的包。。。
画决策树主要需要哪些部分呢：
- 树的各个节点（决策点，即特征）：直接用plt.plot(x,y,…)画点
- 树枝（特征值）：用plt.annotate(‘’,(end),(start),…)来画
- 叶节点（标签）：用plt.plot(x,y,…)来画
明显，这也是一个递归程序： 终止条件：当得到的tree就是label时，画叶节点，终止。 正常：先把tree中顶部的划分特征拿出来，画一个节点，然后取出它的各个特征值，画出各个分支，下面的子树继续丢进递归程序中。
代码：
代码语言：javascript复制plt.figure(figsize=(10,5))
def plot_tree(tree,x,y):
shift = 0.2
if isinstance(tree,str):
plt.plot(x,y,'go')
plt.text(x,y-shift,tree)
return 0
# 先判断是特征还是特征值
feature = list(tree.keys())[0]
assert feature in feature_names,"%s is not a feature!"�ature
plt.plot(x,y,'ro',markersize=10)
plt.text(x,y shift,feature)
for i,v in enumerate(list(tree[feature].keys())): # 取出每一个特征值
next_x = x/3 i
next_y = y - 1
plt.annotate('',(next_x,next_y),(x,y),arrowprops=dict(width=1,shrink=0.05))
plt.text((next_x x)/2,(next_y y)/2,v)
sub_tree = tree[feature][v]
plot_tree(sub_tree,next_x,next_y)
feature_names = ['色泽','根蒂','敲声','纹理','脐部','触感']
plot_tree(t,1,-1)
plt.show()
上面的代码得到的结果是这样的：
诚然，它很丑。甚至树枝都交叉了，但我不想再调了，matplotlib这玩意儿我实在用着别扭。 一个合格的决策树的画法，应该首先计算一下树的基本属性，节点数、深度等等，再好好布局，我没心思继续扣这些细节了，所以虽然我画的决策树丑，但我看着顺眼，看的清楚，就OK！毕竟，写这些代码，主要是为了体验一下其中的思想，真正实战，还是用别人千锤百炼写出来的工具包吧！
五、如何预测？
有了决策树，怎么去预测？ 这个实际上是最简单的了。 给定一个样本，用我们得到的tree字典依次去检查： 先取出tree的第一个划分特征，找到样本对于的特征值，去查询这个特征值对应的子树即可。 代码如下：
代码语言：javascript复制def predict(tree,feature_df):
feature = list(tree.keys())[0]
fv = list(feature_df[feature])[0]
sub_tree = tree[feature][fv]
if isinstance(sub_tree,str):
print(feature ':' fv,'-->',sub_tree,end='')
return sub_tree
else:
print(feature ':' fv,'--> ',end='')
predict(sub_tree,feature_df)
我们拿数据集中的样本一个个检验一下：
代码语言：javascript复制for i in range(1,18):
result = predict(t,df[df["编号"]==i])
print()
输出结果：
代码语言：javascript复制纹理:清晰 --> 根蒂:蜷缩 --> 是
纹理:清晰 --> 根蒂:蜷缩 --> 是
纹理:清晰 --> 根蒂:蜷缩 --> 是
纹理:清晰 --> 根蒂:蜷缩 --> 是
纹理:清晰 --> 根蒂:蜷缩 --> 是
纹理:清晰 --> 根蒂:稍蜷 --> 色泽:青绿 --> 是
纹理:稍糊 --> 触感:软粘 --> 是
纹理:清晰 --> 根蒂:稍蜷 --> 色泽:乌黑 --> 否
纹理:稍糊 --> 触感:硬滑 --> 否
纹理:清晰 --> 根蒂:硬挺 --> 否
纹理:模糊 --> 否
纹理:模糊 --> 否
纹理:稍糊 --> 触感:硬滑 --> 否
纹理:稍糊 --> 触感:硬滑 --> 否
纹理:清晰 --> 根蒂:稍蜷 --> 色泽:乌黑 --> 否
纹理:模糊 --> 否
纹理:稍糊 --> 触感:硬滑 --> 否
发现，除了8号，都预测对了。
那么，我们终于写就了一个简陋的决策树算法啦！终于从坑里爬出来了！
