rpart_var_importance_plot.R

var_imp <- function(rpart_fit) {
  var_imp <- sort(rpart_fit$variable.importance)
  in_model <- levels(rpart_fit$frame$var)[-1]
  
  return(list(var_imp = var_imp, in_model = in_model))
}

var_imp_plot <- function(rpart_fit) {
  vars <- var_imp(rpart_fit)
  groups <- factor(x = rep("less so", length(vars$var_imp)), 
                   levels = c("important", "less so"))
  match_idx <- match(vars$in_model, names(vars$var_imp))
  groups[match_idx] <- "important"
  color <- character(length(vars$var_imp))
  color[match_idx] <- "red"
  color[-match_idx] <- "blue"
  dotchart(vars$var_imp, labels = names(vars$var_imp), cex = 0.7, pch = 19,
           groups = groups, color = color)
}