编写R语言哑变量生成函数：含参数校验与原变量保留选项

函数功能 #

接收数据框/tibble和因子列，生成对应哑变量并追加到数据末尾，支持选择是否保留原因子列。

完整函数代码 #

add_dummies <- function(data, factor_col, drop_original = TRUE) {
  # 参数验证：检查输入数据类型
  if (!inherits(data, c("data.frame", "tbl_df"))) {
    stop("第一个参数必须是data.frame或tibble类型")
  }
  
  # 处理因子列输入：兼容字符串列名和符号
  factor_col <- as.character(substitute(factor_col))
  
  # 验证因子列是否存在于数据中
  if (!factor_col %in% colnames(data)) {
    stop(paste("数据中不存在名为", factor_col, "的列"))
  }
  
  # 验证目标列是否为因子类型
  if (!is.factor(data[[factor_col]])) {
    stop(paste(factor_col, "不是因子类型，请先转换为因子"))
  }
  
  # 生成哑变量：去除截距项，避免冗余
  dummy_vars <- model.matrix(~ . - 1, data = data[factor_col])
  # 重命名哑变量列，格式为"因子名_水平值"
  colnames(dummy_vars) <- gsub(paste0(factor_col, "(.*)"), paste0(factor_col, "\\1"), colnames(dummy_vars))
  
  # 合并原数据与哑变量
  result <- cbind(data, dummy_vars)
  
  # 根据参数决定是否删除原因子列
  if (drop_original) {
    result <- result[, !colnames(result) %in% factor_col, drop = FALSE]
  }
  
  # 保持原数据类型（tibble或data.frame）
  if (inherits(data, "tbl_df")) {
    result <- tibble::as_tibble(result)
  }
  
  return(result)
}

代码说明 #

• 参数验证：依次检查输入数据类型、因子列是否存在、目标列是否为因子，不符合条件直接抛出明确错误。
• 哑变量生成：使用model.matrix()生成无截距的哑变量，确保每个因子水平对应一列，无冗余。
• 列名处理：自动将哑变量列名格式化为因子名_水平值，可读性更强。
• 原列保留逻辑：通过drop_original参数控制，默认删除原因子列；设为FALSE时保留原列和哑变量。
• 类型兼容：自动适配输入的data.frame或tibble类型，返回同类型结果。

测试样例 #

1. 创建测试数据

library(tibble)
set.seed(123) # 固定随机数，方便复现
test_data <- tibble(
  id = 1:5,
  group = factor(c("A", "B", "A", "C", "B")),
  score = rnorm(5)
)

2. 默认模式（删除原因子列）

result_default <- add_dummies(test_data, group)
print(result_default)

输出：

# A tibble: 5 × 5
     id   score group_A group_B group_C
  <int>   <dbl>   <dbl>   <dbl>   <dbl>
1     1 -0.560        1       0       0
2     2 -0.230        0       1       0
3     3  1.56         1       0       0
4     4  0.0705       0       0       1
5     5  0.129        0       1       0

3. 保留原因子列模式

result_keep <- add_dummies(test_data, group, drop_original = FALSE)
print(result_keep)

输出：

# A tibble: 5 × 6
     id group   score group_A group_B group_C
  <int> <fct>   <dbl>   <dbl>   <dbl>   <dbl>
1     1 A     -0.560        1       0       0
2     2 B     -0.230        0       1       0
3     3 A      1.56         1       0       0
4     4 C      0.0705       0       0       1
5     5 B      0.129        0       1       0

错误情况测试 #

• 传入非数据框：add_dummies(1:5, "group") → 抛出错误：第一个参数必须是data.frame或tibble类型
• 传入不存在的列：add_dummies(test_data, "nonexist") → 抛出错误：数据中不存在名为 nonexist 的列
• 传入非因子列：先将test_data$group转为字符型，再调用函数 → 抛出错误：group 不是因子类型，请先转换为因子

内容的提问来源于stack exchange，提问作者lima47