基于sklearn的决策树实现

import numpy as np
import pandas as pd
from sklearn.metrics import confusion_matrix  # 混淆矩阵
from sklearn.model_selection import train_test_split  # 数据分离模块
from sklearn.tree import DecisionTreeClassifier   #  分类决策树
from sklearn.metrics import accuracy_score  # 评价指标
from sklearn.metrics import classification_report   # 生成分类结果报告模块

# 读取数据 importing data
def load_data():
    balance_data = pd.read_csv('https://archive.ics.uci.edu/ml/machine-learning-'+'databases/balance-scale/balance-scale.data',sep=',',header=None)   # 导入数据集，同时设置头部
    print("Dataset Length", len(balance_data))

    print(balance_data.head())
    return balance_data

# 训练集和测试集的分离 splitting the dataset into train and test
def split_dataset(balance_data):

    X = balance_data.values[:, 1:5]  # 提取特征数据
    y = balance_data.values[:, 0]  # 提取数据标签

    X_train, X_test, y_train, y_test=train_test_split(X,y,test_size=0.3,
                                                      random_state=100)  # 进行数据分离

    return X, y, X_train, X_test, y_train, y_test

# 使用基尼系数进行训练 training with giniIndex
def train_using_gini(X_train, y_train):

    # 先建立实例，再进行fit拟合
    clf_gini = DecisionTreeClassifier(criterion="gini"   # 实例化
                                     ,random_state=100
                                     ,max_depth=3
                                     ,min_samples_leaf=5)
    clf_gini.fit(X_train, y_train)  # fit拟合
    return clf_gini

# 使用信息熵进行训练 training with entropy
def train_using_entropy(X_train, y_train):

    # 实例化+fit拟合
    clf_entropy = DecisionTreeClassifier(criterion="entropy"
                                     ,random_state=100
                                     ,max_depth=3
                                     ,min_samples_leaf=5)
    clf_entropy.fit(X_train, y_train)
    return clf_entropy

# 预测功能 make predictions
def prediction(X_test, clf_object):

    y_pred = clf_object.predict(X_test)
    print("Predicted vlaues:")
    print(y_pred)
    return y_pred

# 计算准确率 calculate accuracy
def cal_accuracy(y_test, y_pred):

    print("Confusion Matrix:", confusion_matrix(y_test, y_pred))

    print("Accuracy:", accuracy_score(y_test, y_pred)*100)

    print("Report:", classification_report(y_test, y_pred))

def main():
    data = load_data()
    X, y, X_train, X_test, y_train, y_test = split_dataset(data)
    clf_gini = train_using_gini(X_train, y_train)
    clf_entropy = train_using_entropy(X_train, y_train)

    print("result using gini Index:")
    y_pred_gini = prediction(X_test, clf_gini)
    cal_accuracy(y_test, y_pred_gini)

    print("result using Entropy:")
    y_pred_entropy = prediction(X_test, clf_entropy)
    cal_accuracy(y_test, y_pred_entropy)

if __name__ == "__main__":
    main()