Aquí tienes un ejemplo completo y explicado paso a paso de cómo entrenar y evaluar un árbol de decisión usando scikit-learn con el dataset del Titanic:
🧪 1. Cargar y preparar los datos
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
# Cargar el dataset (puede ser desde seaborn, kaggle o archivo local)
titanic = pd.read_csv('titanic.csv') # o usa seaborn.load_dataset('titanic')
# Selección de variables
features = ['Pclass', 'Sex', 'Age', 'Fare']
target = 'Survived'
# Eliminar nulos simples (sólo para simplificar el ejemplo)
titanic = titanic.dropna(subset=features + [target])
# Convertir variables categóricas a numéricas
titanic = pd.get_dummies(titanic, columns=['Sex'], drop_first=True, dtype=int)
# Variables predictoras y variable objetivo
X = titanic[['Pclass', 'Age', 'Fare', 'Sex_male']]
y = titanic[target]
🧠 2. Dividir en conjunto de entrenamiento y prueba
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
🌲 3. Entrenar el árbol de decisión
model = DecisionTreeClassifier(random_state=42)
model.fit(X_train, y_train)
📊 4. Evaluación del modelo
# Predicciones
y_pred = model.predict(X_test)
# Métricas
print("Accuracy:", accuracy_score(y_test, y_pred))
print("\nClassification Report:\n", classification_report(y_test, y_pred))
print("\nConfusion Matrix:\n", confusion_matrix(y_test, y_pred))
🌳 5. (Opcional) Visualizar el árbol
from sklearn.tree import plot_tree
import matplotlib.pyplot as plt
plt.figure(figsize=(20, 10))
plot_tree(model, feature_names=X.columns, class_names=['No', 'Yes'], filled=True)
plt.show()
✅ Resultado
Con esto obtendrás:
- Un modelo de árbol de decisión entrenado.
- Métricas de precisión (
accuracy, precision, recall, f1-score).
- Visualización clara de cómo el árbol toma decisiones.