机器学习python构建ID3决策树+pygraphviz可视化（代码详解，附注释）

Posted 2021-05-23 nefu_ljw

---

本文由作者nefu_ljw原创，转载请标明出处！

---

本文完整代码和数据集文件可在此下载：https://download.csdn.net/download/ljw_study_in_CSDN/18364420

准备工作

（1）使用西瓜数据集2.0，数据集命名为watermelon2.0.txt。
  注意：该数据集的属性值全部为离散值，无连续值、缺失值。

编号,色泽,根蒂,敲声,纹理,脐部,触感,好瓜
1,青绿,蜷缩,浊响,清晰,凹陷,硬滑,是
2,乌黑,蜷缩,沉闷,清晰,凹陷,硬滑,是
3,乌黑,蜷缩,浊响,清晰,凹陷,硬滑,是
4,青绿,蜷缩,沉闷,清晰,凹陷,硬滑,是
5,浅白,蜷缩,浊响,清晰,凹陷,硬滑,是
6,青绿,稍蜷,浊响,清晰,稍凹,软粘,是
7,乌黑,稍蜷,浊响,稍糊,稍凹,软粘,是
8,乌黑,稍蜷,浊响,清晰,稍凹,硬滑,是
9,乌黑,稍蜷,沉闷,稍糊,稍凹,硬滑,否
10,青绿,硬挺,清脆,清晰,平坦,软粘,否
11,浅白,硬挺,清脆,模糊,平坦,硬滑,否
12,浅白,蜷缩,浊响,模糊,平坦,软粘,否
13,青绿,稍蜷,浊响,稍糊,凹陷,硬滑,否
14,浅白,稍蜷,沉闷,稍糊,凹陷,硬滑,否
15,乌黑,稍蜷,浊响,清晰,稍凹,软粘,否
16,浅白,蜷缩,浊响,模糊,平坦,硬滑,否
17,青绿,蜷缩,沉闷,稍糊,稍凹,硬滑,否

（2）事先安装包：math、numpy、pandas、pygraphviz。

代码实现

由于代码比较长看起来不太方便，下面分别讲解各个模块。

（1）导入所需包

import math
import numpy as np
import pandas as pd
import pygraphviz as pgv

（2）编写决策树分类器，class类型，分为若干个函数模块，整体框架如下（未标明参数和函数体）：

class ljw_DecisionTreeClassifier():

    def __init__(self, ...):
        
    def calculate_entry(self, ...):

    def discrete_Gain(self, ...):

    def root_decision(self, ...):

    def create_tree(self, ...):

    def fit(self, ...):

    def plot_tree(self, ...):

（3）分别编写各个函数。

①：__init__
初始化实例变量G，它是由pygraphviz包生成的图类型。

def __init__(self, G):
    self.G = G

②：calculate_entry
由类别种类及个数计算信息熵。

def calculate_entry(self, C):
    # 传入：只含一列决策属性(类别标签)的数据子集
    # 返回：信息熵
    C_unique, count = np.unique(C, return_counts=True)  # 去重，得到类别标签的所有取值和对应次数
    entry = 0
    count_sum = sum(count)
    for i in range(len(count)):
        p = count[i] / count_sum
        if p != 0:
            entry = entry - p * math.log2(p)
    return entry

③：discrete_Gain
调用calculate_entry函数，计算条件属性的信息增益。

def discrete_Gain(self, subdata):
    # 传入：只含两列的数据子集，包括条件属性和决策属性(类别标签)，其中条件属性的取值必须为离散值
    # 返回：若取该条件属性作为划分点，对应的信息增益
    C = subdata[:, -1]
    E_init = self.calculate_entry(C)  # 划分前信息熵
    E_new = 0  # 划分后信息熵
    A = subdata[:, 0]  # 一列条件属性
    A_unique = np.unique(A)  # 去重，得到条件属性的所有取值
    for i in range(len(A_unique)):  # 遍历条件属性的所有取值
        pos = np.where(A == A_unique[i])  # 条件属性值为A_unique[i]在A中的所有位置(在subdata的位置)
        subdata_son = subdata[pos[0], :]  # 取出条件属性值为A_unique[i]的子集
        C_son = subdata_son[:, -1]
        E_new = E_new + len(subdata_son) / len(subdata) * self.calculate_entry(C_son)
    return E_init - E_new

④：root_decision
调用discrete_Gain函数，选择当前最优划分属性（信息增益最大）作为根节点。

def root_decision(self, df):
    # 传入：含若干条件属性+决策属性的数据子集，dataFrame格式
    # 返回：最优划分属性和对应下标
    data = df.values  # 用.values取出得到numpy封装的ndarray类型数据，可当作"二维矩阵"使用
    col = len(data[0])
    Gain = np.zeros(col - 1)
    for i in range(col - 1):  # 遍历所有条件属性
        subdata = data[:, [i, -1]]  # 取出子矩阵，包括所有行，但是只包括两列：当前条件属性和决策属性
        Gain[i] = self.discrete_Gain(subdata)  # 条件属性取值是离散值
    max_index = np.argmax(Gain)
    final_attribution = df.columns.values[max_index]
    return final_attribution, max_index

⑤：create_tree
这是最核心的函数，递归调用root_decision函数不断生成结点并将其加入到图G。

参照《机器学习》书上的算法递归生成决策树，注意return的三种情况（我写的代码 终止条件3 不是return而是产生叶节点后继续循环，具体详见代码)。

def create_tree(self, df, father_id, edge_name, attribution, A_origin):
    now_id = self.G.number_of_nodes() + 1  # now_id记录递归栈中本层的结点编号
    data = df.values
    C = data[:, -1]
    C_unique, C_count = np.unique(C, return_counts=True)
    # Case 1：所有样本属于同一类别，为叶结点，其标记为该类别，然后返回上层
    if len(C_unique) == 1:
        self.G.add_node(now_id, label=C_unique[0], fontname="Microsoft YaHei", style="filled",
                        color="#B4E7B7")  # 点为绿色(#B4E7B7)，填充(filled)
        self.G.add_edge([father_id, now_id], color="#B4DBFF", penwidth=1.5, fontsize=12,
                        fontname="Microsoft YaHei", label=edge_name)  # 边为蓝色(#B4DBFF)
        return
    A = data[:, :-1]
    A_unique = np.unique(A.astype(str), axis=0)  # 必须用astype变成同一类型才能按行去重
    count_index = np.argmax(C_count)  # 类别值最大出现次数
    C_maxCount_label = C_unique[count_index]
    # Case 2：所有条件属性按行去重后取值相同或为空，为叶结点，其标记为当前出现次数最多的类别，然后返回上层
    if len(A_unique) <= 1:
        self.G.add_node(now_id, label=C_maxCount_label, fontname="Microsoft YaHei", style="filled",
                        color="#B4E7B7")
        self.G.add_edge([father_id, now_id], color="#B4DBFF", penwidth=1.5, fontsize=12,
                        fontname="Microsoft YaHei", label=edge_name)
        return
    # 选择最优划分属性final_attribution
    final_attribution, max_index = self.root_decision(df)
    # 先把自己这个结点加进去
    self.G.add_node(now_id, label=final_attribution, fontname="Microsoft YaHei")  # 点为绿色(#B4E7B7)
    if father_id != -1:
        self.G.add_edge([father_id, now_id], color="#B4DBFF", penwidth=1.5, fontsize=12,
                        fontname="Microsoft YaHei", label=edge_name)  # 边为蓝色(#B4DBFF)
    origin_index = np.where(attribution == final_attribution)[0]
    A_origin_unique = np.unique(A_origin[:, origin_index])
    for i in range(len(A_origin_unique)):  # 遍历最优划分属性在原样本中的所有属性取值
        if not np.isin(A_origin_unique[i], A[:, max_index]):
            # Case 3：原属性取值不在当前属性A中，为叶结点，其标记为当前出现次数最多的类别，然后去下一次循环
            # NOTE：这里我写的代码和书上不同，不是递归进入下层，标记其父结点df中出现次数最多的类别
            # 而是直接把下层的叶结点标记为现在出现次数最多的类别，然后去下一次循环
            next_id = self.G.number_of_nodes() + 1
            self.G.add_node(next_id, label=C_maxCount_label, fontname="Microsoft YaHei", style="filled",
                            color="#B4E7B7")
            self.G.add_edge([now_id, next_id], color="#B4DBFF", penwidth=1.5, fontsize=12,
                            fontname="Microsoft YaHei", label=A_origin_unique[i])
        else:
            df_son = df.loc[df[final_attribution] == A_origin_unique[i]].drop(columns=final_attribution)
            self.create_tree(df_son, now_id, A_origin_unique[i], attribution,
                             A_origin)  # 分支边为A_origin_unique[i]
    return

⑥：fit
训练分类器，得到的模型即图G。

def fit(self, train_df):
    #  train_df要求：datafatame格式(含表头)、最后一列为决策属性、其余列为条件属性
    A_origin = train_df.values[:, :-1]
    attribution = train_df.columns.values
    self.create_tree(train_df, -1, -1, attribution, A_origin)

⑦：plot_tree
将图G可视化，保存为图片文件。

def plot_tree(self, filename):
    self.G.layout()
    self.G.draw(filename, prog="dot")

（4）main函数生成分类器对象，并调用分类器函数。

if __name__ == "__main__":
    # 创建一个分类器对象
    classifier = ljw_DecisionTreeClassifier(G=pgv.AGraph(directed=True))
    # 读取数据
    train_df = pd.read_table("watermelon2.0.txt", sep=",").drop(columns="编号")  # 去掉编号属性
    # 训练得到图G(在分类器中)
    classifier.fit(train_df)
    # 将分类器中的G保存为图片
    classifier.plot_tree("watermelon2.0.png")

输出图像结果

后续可改进优化部分

• 预剪枝与后剪枝
• 处理连续值与缺失值

参考资料

1. 《机器学习》 周志华
2. 决策树算法
3. PyGraphviz (几何图形可视化工具) 简单入门

官方文档：

• NumPy 1.20
• pandas 1.2.4
• PyGraphviz 1.7