one-hot encodes a variable and returns a data.table with new columns

one_hot_encode_var.r

# A helper function to one-hot encode a _single_ variable
# Assumes data.table column or vector input
# Returns a data.table with x columns where x is number of levels
ohev <- function(vars, drop_ref = TRUE) {
  vname <- deparse(substitute(vars))
  lvls <- unlist(unique(vars))
  tmp_list <- list()
  for (lev in lvls) {
    tmp_list[[make.names(paste0(vname,'_',lev))]] <- as.numeric(vars == lev)
  }
  res <- as.data.table(tmp_list)
  if (drop_ref) {
    return(res[,-1])
  } else {
    return(res)
  }
}