Comment tracer un arbre de décision dans R (avec exemple)

En apprentissage automatique , un arbre de décision est un type de modèle qui utilise un ensemble de variables prédictives pour créer un arbre de décision qui prédit la valeur d’une variable de réponse.

Le moyen le plus simple de tracer un arbre de décision dans R consiste à utiliser la fonction prp() du package rpart.plot .

L’exemple suivant montre comment utiliser cette fonction dans la pratique.

Exemple : tracer un arbre de décision dans R

Pour cet exemple, nous utiliserons l’ensemble de données Hitters du package ISLR , qui contient diverses informations sur 263 joueurs de baseball professionnels.

Nous utiliserons cet ensemble de données pour construire un arbre de régression qui utilise les circuits et les années jouées pour prédire le salaire d’un joueur donné.

Le code suivant montre comment ajuster cet arbre de régression et comment utiliser la fonction prp() pour tracer l’arbre :

library(ISLR)
library(rpart)
library(rpart.plot)

#build the initial decision tree
tree <- rpart(Salary ~ Years + HmRun, data=Hitters, control=rpart.control(cp=.0001))

#identify best cp value to use
best <- tree$cptable[which.min(tree$cptable[,"xerror"]),"CP"]

#produce a pruned tree based on the best cp value
pruned_tree <- prune(tree, cp=best)

#plot the pruned tree
prp(pruned_tree)

Notez que nous pouvons également personnaliser l’apparence de l’arbre de décision en utilisant les arguments faclen , extra , roundint et digits dans la fonction prp () :

#plot decision tree using custom arguments
prp(pruned_tree,
    faclen=0, #use full names for factor labels
    extra=1, #display number of observations for each terminal node
    roundint=F, #don't round to integers in output
    digits=5) #display 5 decimal places in output

Nous pouvons voir que l’arbre a six nœuds terminaux.

Chaque nœud terminal affiche le salaire prévu des joueurs de ce nœud ainsi que le nombre d’observations de l’ensemble de données d’origine qui appartiennent à cette note.

Par exemple, nous pouvons voir que dans l’ensemble de données d’origine, il y avait 90 joueurs avec moins de 4,5 ans d’expérience et que leur salaire moyen était de 225,83 000 $ .

Nous pouvons également utiliser l’arbre pour prédire le salaire d’un joueur donné en fonction de ses années d’expérience et de ses circuits moyens.

Par exemple, un joueur qui a 7 ans d’expérience et 4 circuits en moyenne a un salaire prévu de 502,81 000 $ .

C’est l’un des avantages de l’utilisation d’un arbre de décision : nous pouvons facilement visualiser et interpréter les résultats.

Ressources additionnelles

Les didacticiels suivants fournissent des informations supplémentaires sur les arbres de décision :

Une introduction aux arbres de classification et de régression
Arbre de décision vs forêts aléatoires : quelle est la différence ?
Comment ajuster les arbres de classification et de régression dans R