基于R语言扩展Guess Who游戏条件概率计算函数求助

以下是扩展后的函数，既保留原有的筛选功能，又支持计算指定条件下的条件概率：

my_function <- function(dataset, input_var_list = NULL, conditional_var = NULL, var1 = NULL, var2 = NULL, var3 = NULL, var4 = NULL) {
  
  # 从单个变量参数构建条件列表（兼容旧调用方式）
  if (is.null(input_var_list)) {
    input_var_list <- list()
    if (!is.null(var1)) input_var_list$var1 <- var1
    if (!is.null(var2)) input_var_list$var2 <- var2
    if (!is.null(var3)) input_var_list$var3 <- var3
    if (!is.null(var4)) input_var_list$var4 <- var4
  }
  
  # 验证条件列表的有效性
  if (length(input_var_list) > 0) {
    if (!is.list(input_var_list) || is.null(names(input_var_list)) || any(names(input_var_list) == "")) {
      stop("input_var_list必须是命名列表，每个名称对应数据集中的列")
    }
    if (!all(names(input_var_list) %in% colnames(dataset))) {
      missing_vars <- setdiff(names(input_var_list), colnames(dataset))
      stop(paste("数据集中不存在以下变量：", paste(missing_vars, collapse = ", ")))
    }
  }
  
  # 根据条件筛选数据集
  filtered_data <- dataset
  if (length(input_var_list) > 0) {
    for (var_name in names(input_var_list)) {
      filtered_data <- filtered_data[filtered_data[[var_name]] %in% input_var_list[[var_name]], ]
    }
  }
  
  # 若未指定条件概率变量，返回筛选后的数据集（兼容旧功能）
  if (is.null(conditional_var)) {
    return(filtered_data)
  }
  
  # 验证条件概率变量是否存在
  if (!conditional_var %in% colnames(dataset)) {
    stop(paste("条件概率变量", conditional_var, "不存在于数据集中"))
  }
  
  # 计算筛选后数据的总计数
  total_count <- sum(filtered_data$counts)
  
  # 处理总计数为0的情况（避免除以0）
  if (total_count == 0) {
    all_levels <- levels(dataset[[conditional_var]])
    result <- data.frame(
      level = all_levels,
      probability = rep(0, length(all_levels))
    )
    colnames(result)[1] <- conditional_var
    return(result)
  }
  
  # 按条件概率变量分组求和
  counts_per_level <- aggregate(counts ~ ., data = filtered_data[, c(conditional_var, "counts")], sum)
  
  # 获取条件概率变量的所有水平（包括筛选后未出现的）
  all_levels <- data.frame(level = levels(dataset[[conditional_var]]))
  colnames(all_levels)[1] <- conditional_var
  
  # 合并数据，补全缺失水平的计数为0
  result <- merge(all_levels, counts_per_level, by = conditional_var, all.x = TRUE)
  result$counts[is.na(result$counts)] <- 0
  
  # 计算条件概率
  result$probability <- result$counts / total_count
  
  # 返回仅包含条件变量和概率的结果
  return(result[, c(conditional_var, "probability")])
}

使用示例 #

1. 原有筛选功能（兼容旧调用方式） #

筛选黑发/棕发且戴眼镜的运动员：

filtered_data <- my_function(dataset, var1 = c("black", "brown"), var2 = "yes")

2. 计算条件概率 #

计算黑发/棕发且踢足球时，不同眼镜类型的条件概率：

prob_result <- my_function(
  dataset, 
  input_var_list = list(var1 = c("black", "brown"), var3 = "football"), 
  conditional_var = "var2"
)

返回结果示例：

var2 probability
1 contact lenses   0.1500000
2           no   0.3500000
3          yes   0.5000000

功能说明 #

• 兼容性：保留原有的单个变量参数（var1/var2等）调用方式，同时支持新的input_var_list参数批量指定条件
• 鲁棒性：处理筛选后无数据的情况，返回所有水平的概率为0
• 完整性：包含条件概率变量的所有原始水平，即使某些水平在筛选后未出现

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