如何用R语言mice包限制分类变量插补至指定类别子集

要实现仅将分类变量x2的缺失值插补为A或B（排除C），且不影响其他变量的插补，可以通过自定义插补函数修改mice的默认行为，具体步骤如下：

1. 自定义针对x2的插补函数 #

编写一个仅从A、B中抽样的插补函数：用x2中已有的A、B非缺失值拟合二分类回归模型，对缺失值预测属于A的概率，再通过随机抽样生成A/B的插补值。

# 自定义x2的插补函数
impute_x2 <- function(y, ry, x, ...) {
  # 提取x2的非缺失值，仅保留A、B水平用于建模
  y_valid <- y[ry]
  y_valid <- factor(y_valid, levels = c("A", "B"))
  
  # 用预测变量拟合二分类逻辑回归
  fit <- glm(y_valid ~ ., data = data.frame(y_valid, x[ry, ]), family = binomial())
  
  # 对缺失值预测属于A的概率
  prob_A <- predict(fit, newdata = data.frame(x[!ry, ]), type = "response")
  
  # 按概率随机抽样生成A/B的插补值
  imp_values <- ifelse(runif(sum(!ry)) < prob_A, "A", "B")
  
  # 转换回包含C水平的原因子类型，避免因子水平丢失
  factor(imp_values, levels = c("A", "B", "C"))
}

2. 调用mice执行带限制的多重插补 #

指定x2使用自定义插补函数，其他变量保留默认插补方法：

require(mice)
set.seed(12345)
N = 50
miss = sample(0:1, N, prob = c(0.9, 0.1), replace = TRUE)
miss2 = sample(0:1, N, prob = c(0.4, 0.6), replace = TRUE)
df = data.frame(
  x1 = rnorm(N),
  x2 = factor(sample(LETTERS[1:3], N, replace = TRUE)),
  x3 = rnorm(N)
)
df$x2[miss==1] <- NA
df$x3[miss2==1] <- NA

# 定义插补方法：x2用自定义函数，其余变量用默认方法
imp_method <- c("", "impute_x2", "")

# 执行多重插补
imp <- mice(df, m = 2, method = imp_method)

# 查看插补结果
imp1 <- complete(imp, 1)
imp2 <- complete(imp, 2)
dfx2 <- data.frame(df$x2, imp1$x2, imp2$x2)
dfx2[is.na(dfx2$`df.x2`),] 

效果说明 #

运行上述代码后，x2的所有缺失值插补结果只会是A或B，不会出现C；同时原数据中x2的C水平会被保留，其他变量（如x3）的缺失值也能正常完成插补，完全符合需求。

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