Python实现Decision Tree

前言

这篇文章是我用Python对Decision Tree的简单实现，不包含剪枝功能。另外，这个Decision Tree只适用于连续性特征值，离散型的以后有机会再补充。数据集为iris。

1. 导入所需包

import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

2. 加载数据

iris = load_iris()
data = pd.DataFrame(iris.data)
data.columns = iris.feature_names
data['target'] = iris.target

3. 实现Decision Tree

构建树的停止条件有两个：

1. 如果节点的数据集中所有的数据标签相同，那么就将此节点作为叶节点，并将此标签作为判断结果。
2. 如果节点的数据集的特征都用完了，则将此节点作为叶节点，并将当前数据集中最多的标签作为判断结果。

特征选择标准是信息增益。

'''
说明：计算数据集的经验熵

参数：
    labels：数据集的labels
返回值：
    entropy：数据集的经验熵
'''
def cal_entropy(labels):
    data_num = len(labels) # 数据集中数据的个数
    labels_num = {
   } # 用于储存数据集中每个label的个数
    for label in labels:
        if label not in labels_num.keys():
            labels_num[label] = 1
        else:
            labels_num[label] += 1
    entropy = 0.0 # 经验熵
    for key in labels_num:
        prob = float(labels_num[key]) / data_num
        entropy -= prob * np.log2(prob)
    return entropy

'''
说明：计算按照当前特征划分后的经验熵（连续性特征值）

参数：
    data：连续性特征值
    labels：数据对应的labels
返回值：
    min_entropy：最小经验熵
    best_partition：使得经验熵最小的划分点
'''

def entropy_for_continuous(data, labels):
    sorted_data = data.sort_values().reset_index(drop=True) # 将数据排序
    partitions = [(sorted_data[i] + sorted_data[i + 1]) / 2 for i in range(0, len(sorted_data) - 1)] # 计算相邻两个数字的中位数
    entropy_list = []
    # 计算按照不同划分点划分后的经验熵
    for partition in partitions:
        smaller_group = data[data <= partition]