如何从rpart提取含替代分割的决策树规则并转为可调整R代码

一、summary输出中路径符号的含义 #

summary(mytree)输出里的RL--这类字符串，每个字符对应决策树一层内部节点的分支选择：

• L：当前层使用主分割变量的条件，走左分支
• R：当前层使用主分割变量的条件，走右分支
• -：当前层主分割变量存在缺失，使用替代分割变量的条件决定分支（或该节点为叶子节点，无主分割）

以RL--为例：

1. 第一层节点：用主分割变量判断，走右分支
2. 第二层节点：用主分割变量判断，走左分支
3. 第三、四层节点：主分割变量缺失，依赖替代分割完成分支

二、从rpart对象提取替代分割逻辑 #

rpart对象内置了所有分割信息，核心字段：

• mytree$splits：所有主分割的变量、阈值、分支方向
• mytree$surrogate：每个内部节点的替代分割列表（含变量、阈值、匹配率）
• mytree$frame：节点类型（内部/叶子）、预测类别等信息
• mytree$where：每个样本最终归属的叶子节点ID

提取单节点的主+替代分割规则 #

以下函数可输入节点ID，返回该节点的主分割+所有替代分割的条件：

extract_split_rules <- function(tree, node_id) {
  # 获取主分割
  main_split_idx <- which(rownames(tree$splits) == tree$frame$var[node_id])
  if (length(main_split_idx) == 0 || tree$frame$var[node_id] == "<leaf>") return(NULL)
  
  main_var <- rownames(tree$splits)[main_split_idx]
  main_thresh <- tree$splits[main_split_idx, "index"]
  # 处理分类变量分割（示例中为"+"/"-"类型）
  main_condition <- if (tree$variable[main_var]$type == "categorical") {
    sprintf("%s == '%s'", main_var, names(tree$splits[index == main_thresh, "index"]))
  } else {
    sprintf("%s <= %f", main_var, main_thresh)
  }
  
  # 获取替代分割
  surrogate_list <- tree$surrogate[[node_id]]
  if (is.null(surrogate_list)) return(list(main = main_condition))
  
  surrogate_conditions <- lapply(seq(nrow(surrogate_list)), function(i) {
    surr_var <- rownames(surrogate_list)[i]
    surr_thresh <- surrogate_list[i, "index"]
    surr_dir <- surrogate_list[i, "direction"]
    # 方向1对应主分割左分支，-1对应右分支
    base_cond <- if (tree$variable[surr_var]$type == "categorical") {
      sprintf("%s == '%s'", surr_var, names(tree$splits[index == surr_thresh, "index"]))
    } else {
      if (surr_dir == 1) sprintf("%s <= %f", surr_var, surr_thresh) else sprintf("%s > %f", surr_var, surr_thresh)
    }
    # 替代分割仅在主变量缺失时生效
    sprintf("%s & is.na(%s)", base_cond, main_var)
  })
  
  list(main = main_condition, surrogates = surrogate_conditions)
}

三、生成dplyr::case_when格式的代码 #

遍历所有叶子节点，将每个叶子的路径规则（主分割+替代分割串联）转换为case_when格式：

library(dplyr)

generate_case_when <- function(tree) {
  # 获取所有叶子节点ID
  leaf_nodes <- rownames(tree$frame[tree$frame$var == "<leaf>", ])
  
  # 生成每个叶子的完整规则
  leaf_rules <- lapply(leaf_nodes, function(node) {
    path_conds <- character(0)
    current_node <- as.integer(node)
    
    # 从叶子回溯到根节点，拼接路径条件
    while (current_node != 1) {
      parent_node <- floor(current_node / 2)
      is_left_child <- (current_node %% 2) == 0
      split_rules <- extract_split_rules(tree, parent_node)
      
      if (!is.null(split_rules)) {
        # 左分支用原条件，右分支取反
        main_cond <- if (is_left_child) split_rules$main else sprintf("!(%s)", split_rules$main)
        surr_conds <- lapply(split_rules$surrogates, function(surr) {
          if (is_left_child) surr else sprintf("!(%s)", gsub("is.na\\((.*?)\\)", "is.na(\\1)", surr))
        })
        # 主分割和替代分割是"或"的关系
        node_cond <- paste(c(main_cond, surr_conds), collapse = " | ")
        path_conds <- c(node_cond, path_conds)
      }
      current_node <- parent_node
    }
    
    # 路径上的所有条件是"与"的关系
    full_cond <- paste(path_conds, collapse = " & ")
    pred_class <- as.character(tree$frame$yval[tree$frame$var == "<leaf>"][which(rownames(tree$frame) == node)])
    list(condition = full_cond, class = pred_class)
  })
  
  # 拼接成case_when代码字符串
  case_lines <- sapply(leaf_rules, function(rule) {
    sprintf("  %s ~ '%s',", rule$condition, rule$class)
  })
  case_code <- paste0("dplyr::case_when(\n", paste(case_lines, collapse = "\n"), "\n  TRUE ~ 'Unknown'\n)")
  cat(case_code)
}

# 运行生成代码
generate_case_when(mytree)

四、手动调整分支的简化方法 #

若需手动修改分支，可将rpart树转为partykit对象，其规则提取和结构修改更直观：

library(partykit)
# 转换rpart树为partykit对象
tree_party <- as.party(mytree)
# 查看带替代分割的完整规则
print(tree_party, type = "extended")
# 手动修改分支（示例：修改节点2的分割规则）
tree_party[[2]]$split <- partykit::split_node(varid = 2, breaks = "-")
# 可直接使用修改后的树预测，或重新生成规则

内容的提问来源于stack exchange，提问作者cliffhanger-be