如何基于底层原理改写R语言中添加公式交互项的函数？

我需要编写一个R函数，用于给formula添加交互项。比如，传入公式mpg ~ cyl + gear + disp、处理变量cyl和控制变量字符向量c("gear","disp")，返回mpg ~ cyl + cyl * gear + cyl * disp。

理想状态下，如果某个控制变量不在原公式中，或者对应的交互项已经存在，函数要返回错误提示。

我目前写了一个版本能运行，但用的是字符串操作，容易出错且效率低，想知道怎么基于公式的底层原理来改写这个函数？

现有字符串实现代码 #

#' 为公式添加交互项
#' 
#' @param form 一个formula对象
#' @param treat 处理变量（字符串）
#' @param controls 控制变量的字符向量
#' @return 已添加`treat`与每个`controls`变量间交互项的formula
#' @export
#' @examples
#' reformulas_addints(mpg ~ cyl + gear, "cyl", c("gear"))
#' reformulas_addints(mpg ~ cyl + gear + disp, "cyl", c("gear", "disp"))
#' reformulas_addints(mpg ~ cyl + gear, "cyl", c("gears"))
#' reformulas_addints(mpg ~ cyl + cyl*gear, "cyl", c("gear"))
reformulas_addints <- function(form, treat, controls) {
  form_str <- as.character(form)
  for (control in controls) {
    if(!stringr::str_detect(form_str, control)){
      stop(paste0("The variable '", control, "' is not in the formula."))
    }
    patt <- paste0(r"(\s*)",treat,r"(\s*\*\s*)",control, r"(\s*)")
    if(stringr::str_detect(form_str, patt)){
      stop(paste0("The interaction '",treat, " * ", control, "' is already in the formula."))
    }
    form_str <- stringr::str_replace(
      form_str,
      paste0("\\b", control, "\\b"),
      paste0(treat, " * ", control)
    )
  }
  return(as.formula(form_str))
}

预期输出示例 #

# 预期输出
reformulas_addints(mpg ~ cyl + gear, "cyl", c("gear"))
# mpg ~ cyl + cyl * gear
# 也可接受以下形式
# mpg ~ cyl + gear + cyl:gear
reformulas_addints(mpg ~ cyl + gear + disp, "cyl", c("gear", "disp"))
# mpg ~ cyl + cyl * gear + cyl * disp
# 也可接受以下形式
# mpg ~ cyl + gear + disp + cyl:gear + cyl:disp
reformulas_addints(mpg ~ cyl + gear + disp + hp, "cyl", c("gear", "disp"))
# mpg ~ cyl + cyl * gear + cyl * disp + hp
# 也可接受以下形式
# mpg ~ cyl + gear + disp + cyl:gear + cyl:disp + hp
# 注意：`hp`未参与交互
reformulas_addints(mpg ~ cyl + gear, "cyl", c("gears"))
# 错误提示：The variable 'gears' is not in the formula.
reformulas_addints(mpg ~ cyl + cyl*gear, "cyl", c("gear"))
# 错误提示：The interaction 'cyl * gear' is already in the formula.

基于公式底层原理的改写方案 #

利用R的terms对象解析公式结构，避免字符串匹配的误差，实现更可靠的逻辑：

#' 为公式添加交互项（底层实现版）
#' 
#' @param form 一个formula对象
#' @param treat 处理变量（字符串）
#' @param controls 控制变量的字符向量
#' @return 已添加`treat`与每个`controls`变量间交互项的formula
#' @export
#' @examples
#' add_interactions(mpg ~ cyl + gear, "cyl", c("gear"))
#' add_interactions(mpg ~ cyl + gear + disp, "cyl", c("gear", "disp"))
#' add_interactions(mpg ~ cyl + gear, "cyl", c("gears"))
#' add_interactions(mpg ~ cyl + cyl*gear, "cyl", c("gear"))
add_interactions <- function(form, treat, controls) {
  # 解析公式为结构化terms对象
  form_terms <- terms(form)
  # 获取公式中所有项的标签（含交互项）
  all_terms <- attr(form_terms, "term.labels")
  # 提取所有单个变量（拆分交互项）
  single_vars <- unique(unlist(strsplit(all_terms, ":|\\*")))
  
  # 检查控制变量是否存在
  missing_controls <- setdiff(controls, single_vars)
  if (length(missing_controls) > 0) {
    stop(sprintf("变量'%s'不在公式中", paste(missing_controls, collapse = "', '")))
  }
  
  # 检查交互项是否已存在（terms中交互项以treat:control形式存储）
  target_interactions <- sapply(controls, function(ctrl) paste(treat, ctrl, sep = ":"))
  existing_interactions <- intersect(target_interactions, all_terms)
  if (length(existing_interactions) > 0) {
    int_str <- gsub(":", " * ", existing_interactions)
    stop(sprintf("交互项'%s'已存在于公式中", paste(int_str, collapse = "', '")))
  }
  
  # 构造新的项列表：替换单个控制变量为treat*control，其余项保留
  new_term_list <- lapply(all_terms, function(term) {
    if (term %in% controls) paste(treat, term, sep = " * ") else term
  })
  
  # 拼接新公式
  new_form <- as.formula(
    paste(deparse(form[[2]]), "~", paste(unlist(new_term_list), collapse = " + "))
  )
  return(new_form)
}

核心逻辑说明 #

1. 结构化解析公式：通过terms()把公式拆解为可直接操作的项标签，避免字符串匹配的误判（比如变量名包含其他变量子串的情况）。
2. 变量存在性校验：从拆分后的单个变量列表中检查控制变量是否存在，结果更准确。
3. 交互项存在性校验：直接检查treat:control形式的交互项（treat*control会被terms解析为主效应+交互项的组合，交互项以:形式存储）。
4. 构造新公式：遍历原公式的项，仅替换需要添加交互的控制变量，其余项完整保留，最后拼接为合法的formula对象。

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