决策树实战：用Python手写ID3算法（附完整代码与数据集）

weixin_30879169

369人浏览 · 2026-03-26 03:04:26

weixin_30879169 · 2026-03-26 03:04:26 发布

决策树实战：用Python手写ID3算法（附完整代码与数据集）

决策树作为机器学习中最直观的算法之一，其核心思想就像我们日常做决策时的思考过程——通过一系列的判断条件逐步缩小范围，最终得出结论。今天我们将从零开始实现经典的ID3决策树算法，通过代码深入理解信息增益的计算、递归构建树结构等关键环节。

1. 环境准备与数据加载

在开始之前，我们需要准备Python环境和演示数据集。这里我们使用一个经典的"贷款审批"数据集，包含年龄、收入、信用等级等特征，以及最终的审批结果。

import numpy as np
import pandas as pd
from math import log2

# 创建演示数据集
data = {
    '年龄': ['青年', '青年', '青年', '青年', '青年', '中年', '中年', '中年', '中年', '中年', '老年', '老年', '老年', '老年', '老年'],
    '有工作': ['否', '否', '是', '是', '否', '否', '否', '是', '否', '否', '否', '否', '是', '是', '否'],
    '有房子': ['否', '否', '否', '是', '否', '否', '否', '是', '是', '是', '是', '是', '否', '否', '否'],
    '信用': ['一般', '好', '好', '一般', '一般', '一般', '好', '好', '非常好', '非常好', '非常好', '好', '好', '非常好', '一般'],
    '类别': ['否', '否', '是', '是', '否', '否', '否', '是', '是', '是', '是', '是', '是', '是', '否']
}

df = pd.DataFrame(data)
print("数据集预览:")
print(df.head())

这个数据集包含了15条记录，每条记录有4个特征和1个类别标签。我们的目标是构建一个决策树模型，能够根据这些特征预测贷款是否会被批准。

2. 核心概念：信息熵与信息增益

ID3算法的核心是信息增益，而要理解信息增益，首先需要掌握信息熵的概念。

信息熵是度量数据集不确定性的指标，计算公式为：

def entropy(data):
    """计算数据集的信息熵"""
    labels = data.iloc[:, -1]
    total = len(labels)
    counts = labels.value_counts()
    ent = 0.0
    for count in counts:
        prob = count / total
        ent -= prob * log2(prob)
    return ent

信息增益则表示使用某个特征划分数据集后，不确定性减少的程度：

def info_gain(data, feature):
    """计算指定特征的信息增益"""
    base_entropy = entropy(data)
    feature_values = data[feature].unique()
    weighted_entropy = 0.0
    
    for value in feature_values:
        subset = data[data[feature] == value]
        weight = len(subset) / len(data)
        weighted_entropy += weight * entropy(subset)
    
    return base_entropy - weighted_entropy

让我们计算各特征的信息增益：

features = df.columns[:-1]
for feature in features:
    print(f"{feature}的信息增益: {info_gain(df, feature):.4f}")

输出结果将显示哪个特征对分类最有帮助，这正是ID3算法选择划分特征的依据。

3. 递归构建决策树

有了信息增益的计算方法，我们就可以递归地构建决策树了。决策树的每个节点将存储以下信息：

如果是叶节点：存储类别标签
如果是内部节点：存储划分特征和子节点

class DecisionNode:
    """决策树节点类"""
    def __init__(self, feature=None, children=None, label=None):
        self.feature = feature   # 划分特征
        self.children = children or {}  # 子节点字典 {特征值: 子节点}
        self.label = label       # 叶节点的类别标签

def build_tree(data, features):
    """递归构建决策树"""
    # 如果所有样本属于同一类别，返回叶节点
    if len(data['类别'].unique()) == 1:
        return DecisionNode(label=data['类别'].iloc[0])
    
    # 如果没有特征可用，返回样本数最多的类别
    if not features:
        majority_label = data['类别'].mode()[0]
        return DecisionNode(label=majority_label)
    
    # 选择最佳划分特征
    best_feature = max(features, key=lambda f: info_gain(data, f))
    
    # 递归构建子树
    children = {}
    for value in data[best_feature].unique():
        subset = data[data[best_feature] == value]
        if len(subset) == 0:
            majority_label = data['类别'].mode()[0]
            children[value] = DecisionNode(label=majority_label)
        else:
            remaining_features = [f for f in features if f != best_feature]
            children[value] = build_tree(subset, remaining_features)
    
    return DecisionNode(feature=best_feature, children=children)

构建决策树并打印结构：

tree = build_tree(df, list(features))

def print_tree(node, indent=""):
    """打印决策树结构"""
    if node.label is not None:
        print(indent + "类别:", node.label)
    else:
        print(indent + "特征:", node.feature)
        for value, child in node.children.items():
            print(indent + f"  {value}:")
            print_tree(child, indent + "    ")

print("决策树结构:")
print_tree(tree)

4. 决策树的可视化与预测

为了更直观地理解决策树，我们可以使用graphviz库进行可视化：

from graphviz import Digraph

def visualize_tree(tree):
    """可视化决策树"""
    dot = Digraph()
    node_id = 0
    
    def add_nodes(node, parent_id=None, edge_label=None):
        nonlocal node_id
        current_id = str(node_id)
        node_id += 1
        
        if node.label is not None:
            dot.node(current_id, label=f"类别: {node.label}", shape='box')
        else:
            dot.node(current_id, label=f"特征: {node.feature}")
            
        if parent_id is not None:
            dot.edge(parent_id, current_id, label=edge_label)
            
        if node.children:
            for value, child in node.children.items():
                add_nodes(child, current_id, str(value))
    
    add_nodes(tree)
    return dot

# 显示决策树
visualize_tree(tree)

有了构建好的决策树，我们就可以用它来进行预测了：

def predict(tree, sample):
    """使用决策树进行预测"""
    if tree.label is not None:
        return tree.label
    
    feature_value = sample[tree.feature]
    if feature_value not in tree.children:
        return None  # 未知特征值
    
    return predict(tree.children[feature_value], sample)

# 测试预测
test_sample = {'年龄': '中年', '有工作': '是', '有房子': '否', '信用': '好'}
print("预测结果:", predict(tree, test_sample))

5. 算法优化与扩展

虽然我们实现了基本的ID3算法，但在实际应用中还需要考虑以下优化：

连续值处理：ID3只能处理离散特征，可以通过二分法将连续特征离散化
缺失值处理：可以采用概率分配或最常见值填充等方法
剪枝策略：预剪枝和后剪枝可以防止过拟合
其他划分准则：如C4.5的信息增益率、CART的基尼指数

这里我们实现一个简单的预剪枝版本：

def build_tree_with_pruning(data, features, max_depth=3, min_samples=2):
    """带预剪枝的决策树构建"""
    if len(data['类别'].unique()) == 1:
        return DecisionNode(label=data['类别'].iloc[0])
    
    if not features or len(data) < min_samples or max_depth <= 0:
        majority_label = data['类别'].mode()[0]
        return DecisionNode(label=majority_label)
    
    best_feature = max(features, key=lambda f: info_gain(data, f))
    
    children = {}
    for value in data[best_feature].unique():
        subset = data[data[best_feature] == value]
        remaining_features = [f for f in features if f != best_feature]
        children[value] = build_tree_with_pruning(
            subset, remaining_features, max_depth-1, min_samples)
    
    return DecisionNode(feature=best_feature, children=children)

# 构建带预剪枝的决策树
pruned_tree = build_tree_with_pruning(df, list(features), max_depth=2)
print("剪枝后的决策树:")
print_tree(pruned_tree)

决策树虽然直观易懂，但也有其局限性。在实际项目中，我们通常会使用随机森林或梯度提升树等集成方法，它们基于决策树但能提供更好的泛化性能。