Colton Star Blog

1
import pandas as pd
2
from sklearn.ensemble import RandomForestClassifier
3
data_path= "train.csv"
4
train_data = pd.read_csv(data_path)
5
test_data = pd.read_csv("test.csv")
6
y = train_data["Survived"]
7
features = ["Pclass", "Sex", "SibSp", "Parch"]
8
X = pd.get_dummies(train_data[features])
9
X_test = pd.get_dummies(test_data[features])
10
model = RandomForestClassifier(
11
    n_estimators=100,        # 采用100棵决策树
12
    max_depth=10,            # 树的深度为10，太高容易过拟合
13
    random_state=2,
14
    n_jobs=-1,               # 使用所有CPU核心并行
15
    verbose=1                # 显示训练进度
16
)
17
model.fit(X, y)
18
predictions = model.predict(X_test)
19
output = pd.DataFrame({'PassengerId': test_data.PassengerId, 'Survived': predictions})
20
output.to_csv('submission.csv', index=False)
21
print("Your submission was successfully saved")