决策树的python实现

以下为决策树的python实现，决策树算法可选ID3、C4.5、CART。

import pandas as pd
import numpy as np


def entropy(data):
    """
    计算数据集的熵。
    """
    n = len(data)  # 计算数据集的长度
    labels = data.iloc[:, -1]  # 获取数据集的标签列
    label_counts = labels.value_counts()  # 统计每个标签的出现次数
    entropy = 0
    # 计算熵
    for label in label_counts.index:
        p = label_counts[label] / n
        entropy -= p * np.log2(p)
    return entropy


def information_gain(data, feature_col):
    """
    计算特征的信息增益。
    :param data: 包含特征和标签的DataFrame。
    :param feature_col: 特征列的索引。
    :return: 特征的信息增益。
    """
    n = len(data)
    feature_values = data.iloc[:, feature_col]  # 获取特征列的值
    feature_counts = feature_values.value_counts()  # 计算特征值的频数
    entropy_before = entropy(data)  # 计算分割前的熵
    entropy_after = 0
    # 计算分割后的熵
    for value in feature_counts.index:
        subset = data[data.iloc[:, feature_col] == value]  # 获取特征值对应的子集
        p = len(subset) / n  # 计算子集在总数据集中的比例
        entropy_after += p * entropy(subset)  # 累加分割后的熵
    return entropy_before - entropy_after  # 返回信息增益


def information_gain_ratio(data, feature_col):
    """
    计算特征的信息增益比。
    """
    return information_gain(data, feature_col) / entropy(data.iloc[:, feature_col])


def gini_index(data):
    """
    计算数据集的基尼指数。
    """
    n = len(data)
    labels = data.iloc[:, -1]
    label_counts = labels.value_counts()
    gini = 1
    # 计算基尼指数
    for label in label_counts.index:
        p = label_counts[label] / n
        gini -= p ** 2
    return gini


def find_best_split(data, algorithm='ID3'):
    """
    找到最佳分割特征。
    :param data: 包含特征和标签的DataFrame。
    :param algorithm: 算法类型，可选ID3、C4.5、CART。
    :return: 最佳分割特征的索引。
    """
    algorithms = {'ID3': information_gain, 'C4.5': information_gain_ratio, 'CART': gini_index}  # 定义不同算法的增益计算函数
    best_gain = 0
    best_feature = None
    # 遍历所有特征列，找到信息增益或信息增益比或基尼指数最大的特征
    for feature_col in range(data.shape[1] - 1):
        gain = algorithms[algorithm](data, feature_col)
        if gain > best_gain:
            best_gain = gain
            best_feature = feature_col
    return best_feature


def build_tree(data, depth=0):
    """
    构建决策树。
    :param data: 包含特征和标签的DataFrame。
    :param depth: 决策树的深度。
    :return: 决策树的字典表示。
    """
    labels = data.iloc[:, -1]
    label_counts = labels.value_counts()
    if len(label_counts) == 1 or depth >= 10: return label_counts.index[0]  # 如果所有标签相同或达到最大深度，返回第一个标签
    feature_col = find_best_split(data)  # 找到最佳分割特征列
    if feature_col is None: return label_counts.index[0]  # 如果没有找到最佳分割特征列，返回第一个标签

    tree = {}
    feature_values = data.iloc[:, feature_col]  # 获取特征列的值
    # 遍历特征列的唯一值，递归构建子树
    for value in feature_values.unique():
        subset = data[data.iloc[:, feature_col] == value]  # 获取特征值对应的子集
        subtree = build_tree(subset, depth + 1)  # 递归构建子树
        tree[value] = subtree  # 构建子树的字典表示
    return {data.columns[feature_col]: tree}  # 返回决策树的字典表示


def classify(tree, row):
    """
    使用决策树进行分类。
    :param tree: 字典表示的决策树。
    :param row: 待分类的样本。
    :return: 分类结果。
    """
    if isinstance(tree, dict):  # 检查当前节点是否为决策树节点
        for key, subtree in tree.items():  # 遍历当前节点的子节点
            if row[key] in subtree:  # 检查样本属性是否在子节点中
                return classify(subtree[row[key]], row)  # 递归分类
    else: return tree  # 返回分类结果


if __name__ == '__main__':
    # 示例数据集,包含4个特征和1个标签
    dataset = pd.DataFrame([[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']], columns=['F1-AGE', 'F2-WORK', 'F3-HOME', 'F4-LOAN', 'F5-CLASS'])
    tree = build_tree(dataset)  # 构建决策树
    print('决策树:', tree)
    # 遍历数据集并使用决策树进行分类
    for i, row in dataset.iterrows():
        print(row[-1], classify(tree, row), end='   ')

输出：

决策树: {'F3-HOME': {0: {'F2-WORK': {0: 'no', 1: 'yes'}}, 1: 'yes'}}
no no   no no   yes yes   yes yes   no no   no no   no no   yes yes   yes yes   yes yes   yes yes   yes yes   yes yes   yes yes   no no