from sklearn.model_selection import train_test_split
import pandas as pd
from sklearn.feature_extraction import DictVectorizer
from sklearn.tree import DecisionTreeClassifier,export_graphviz
def decision_tree():
    #获取数据
    data = pd.read_csv("http://biostat.mc.vanderbilt.edu/wiki/pub/Main/DataSets/titanic.txt")
    #特征值和目标值
    x = data[['pclass', 'age', 'sex', 'name', 'room']]
    y = data['survived']
    #处理缺失值
    x['age'].fillna(x['age'].mean(), inplace=True)#fillna(替换的值，inplace表示替换默认为False）
    #划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25)
    #特征工程（类别信息进行ONE-HOT编码）
    Dict = DictVectorizer(sparse=False)#对字典数据进行特征值化,不按照sparse矩阵展示
    x_train = Dict.fit_transform(x_train.to_dict(orient="records"))#to_dict将数据转换成字典，orient="record"默认按照行进行
    x_test = Dict.fit_transform(x_test.to_dict(orient="records"))
    #决策树预测
    dt = DecisionTreeClassifier(max_depth=5)
    dt.fit(x_train, y_train)
    #准确率评价
    print("准确率为", dt.score(x_test, y_test))
    #导出决策树结构
    export_graphviz(dt, out_file='./tree.dot', feature_names=['pclass', 'age', 'sex', 'name', 'room'])
    return None
if __name__== "__main__":
    decision_tree()