IIS_2023_1/shadaev_anton_lab_3/titanic.py
2023-11-03 15:50:54 +04:00

41 lines
1.5 KiB
Python

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.preprocessing import LabelEncoder
# Загрузка данных
titanic_data = pd.read_csv("titanic_ds.csv")
# Обработка пропущенных значений
titanic_data['Embarked'].fillna('S', inplace=True)
encoder = LabelEncoder()
# Преобразование строковых данных в числовые
titanic_data['Pclass'] = encoder.fit_transform(titanic_data['Pclass'])
titanic_data['Cabin'] = encoder.fit_transform(titanic_data['Cabin'])
titanic_data['Embarked'] = encoder.fit_transform(titanic_data['Embarked'])
# Вывод первых 10 записей
print(titanic_data[['PassengerId', 'Pclass', 'Cabin', 'Embarked']].head(10))
# Разделение данных на признаки и целевую переменную
X = titanic_data[['Pclass', 'Cabin', 'Embarked']]
y = titanic_data['Survived']
# Разделение данных на обучающий и тестовый наборы
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1)
# Создание модели
model = DecisionTreeClassifier()
# Обучение модели
model.fit(X_train, y_train)
# Вычисление важности признаков
importances = model.feature_importances_
# Вывод важности признаков
for i, importance in enumerate(importances):
print(f'Признак {i}: {importance}')