#!/usr/bin/env python
# coding: utf-8

# In[28]:


import pandas as pd


# In[29]:


# 1、获取数据
titianic = pd.read_csv("./titanic.csv")
titianic.head()


# In[30]:


# 筛选特征值和目标值
x = titianic[["pclass", "age", "sex"]]
y = titianic["survived"]


# In[31]:


x.head()


# In[32]:


y.head()


# In[33]:


# 2、数据处理
# 1) 缺失值处理 age字段
x.loc[:, "age"] = x["age"].fillna(x["age"].mean())


# In[35]:


# 2) 转换成字典
x = x.to_dict(orient="records")


# In[36]:


from sklearn.model_selection import train_test_split

# 3、数据集划分
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=22)


# In[37]:


# 4、字典特征抽取
from sklearn.feature_extraction import DictVectorizer


# In[38]:


transfer = DictVectorizer()
x_train = transfer.fit_transform(x_train)
x_test = transfer.transform(x_test)


# In[39]:


from sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz

# 5、模型训练和预测
estimator = DecisionTreeClassifier(criterion="entropy")
estimator.fit(x_train, y_train)
# 6、 模型评估
# 方法一: 直接比对真实值和预测值
y_predict = estimator.predict(x_test)
print("直接比对真实值和预测值:\n", y_test == y_predict)

# 方法二: 计算准确率
score = estimator.score(x_test, y_test)
print("准确率: \n", score)

# 6) 决策树的可视化
export_graphviz(estimator, out_file="./titian_tree.dot", feature_names=transfer.get_feature_names_out())