R/replace-na.R

In camroach87/myhelpr: Helper Functions

#' Replace NA values
#'
#' Replaces NA values with a defined alternative. For a numeric vector, this can
#' be either the mean or median of the vector. For a character or factor vector,
#' this will be the mode.
#'
#' @param x Vector that contains NA values.
#' @param method Replacement method for NA values. For a numeric vector the user
#'   can choose between `mean` and `median`. For character or factor vectors the
#'   mode is automatically chosen.
#'
#' @return A vector without NA values.
#' @export
#'
#'
#' @examples
#' x <- c(1, 3, 5, 5, 6, NA, 10)
#' replace_na(x)
#' # 1  3  5  5  6  5 10
#'
#' x <- c("Cat", "Dog", "Catdog", NA, "Cat")
#' replace_na(x)
#' # "Cat"    "Dog"    "Catdog" "Cat"    "Cat"
replace_na <- function(x, method) {
  UseMethod("replace_na")
}

#' @export
replace_na.default <- function(x, method) {
  cat("Error: Incorrect format for vector. Function requires either numeric,",
      "character or factor input.")
}

#' @export
replace_na.numeric <- function(x, method = "mean") {
  if (method == "mean") {
    x[is.na(x)] <- mean(x, na.rm = TRUE)
  } else if (method == "median") {
    x[is.na(x)] <- stats::median(x, na.rm = TRUE)
  }
  x
}

#' @export
replace_na.character <- function(x, method) {
  x_mode <- tibble::data_frame(x = x) %>%
    dplyr::count(x) %>%
    dplyr::arrange(desc(n)) %>%
    dplyr::slice(1) %>%
    .$x
  x[is.na(x)] <- x_mode
  x
}

#' @export
replace_na.factor <- function(x, method) {
  replace_na(as.character(x))
}