按ID条件填充NA：根据A/B值填充不同前置NA数量（R语言）

问题背景与数据示例 #

现有包含分组id列和带缺失值（NA）的x列的data.frame：

x <- c(NA, NA, NA, NA, "A", NA, NA, "B", "A", NA, NA, NA, NA, "B", NA, "A")  
id <- rep(1:3, c(8, 6, 2))
df <- data.frame(id = id, x = x)

填充需求 #

按id分组处理x列的缺失值：

• 当x的非缺失值为"A"时，将该值前2个连续NA替换为"A"
• 当x的非缺失值为"B"时，将该值前3个连续NA替换为"B"

预期输出 #

处理后的结果如下：

#    id    x
# 1   1 <NA>
# 2   1 <NA>
# 3   1    A
# 4   1    A
# 5   1    A
# 6   1    B
# 7   1    B
# 8   1    B
# 9   2    A
# 10  2 <NA>
# 11  2    B
# 12  2    B
# 13  2    B
# 14  2    B
# 15  3    A
# 16  3    A

高效解决方案 #

1. 基于data.table的实现（适配超大规模数据） #

data.table的分组操作和向量化处理适合百万级数据集，核心思路是先标记每个非缺失值的填充范围，再反向填充：

library(data.table)
setDT(df)

# 定义每个值对应的填充步数
fill_steps <- c(A = 2, B = 3)

df[, {
  # 获取组内非缺失值的位置和对应值
  non_na_pos <- which(!is.na(x))
  non_na_vals <- x[non_na_pos]
  # 生成每个非缺失值需要填充的前序位置范围
  fill_pos_list <- mapply(function(pos, val) {
    max(1, pos - fill_steps[val]):(pos - 1)
  }, non_na_pos, non_na_vals, SIMPLIFY = FALSE)
  # 合并填充位置与对应值，过滤超出组范围的位置
  fill_pos <- unlist(fill_pos_list)
  fill_vals <- rep(non_na_vals, fill_steps[non_na_vals])
  valid_idx <- fill_pos <= .N
  fill_pos <- fill_pos[valid_idx]
  fill_vals <- fill_vals[valid_idx]
  # 赋值填充，保留原始非缺失值
  x[fill_pos] <- fill_vals
  .SD
}, by = id]

2. 基于collapse的实现（更轻量高效） #

collapse包针对分组操作做了底层优化，速度更快、内存占用更低，适合10万组+800万行的超大数据集：

library(collapse)

# 定义组内填充函数
fill_func <- function(x) {
  fill_steps <- c(A = 2, B = 3)
  non_na <- which(!is.na(x))
  if(length(non_na) == 0) return(x)
  # 遍历每个非缺失值，向前填充对应步数的NA
  for(i in seq_along(non_na)) {
    val <- x[non_na[i]]
    start <- max(1, non_na[i] - fill_steps[val])
    end <- non_na[i] - 1
    if(start <= end) x[start:end] <- val
  }
  x
}

# 按id分组应用填充函数
df <- ftransform(df, x = fby(x, id, fill_func))

方案说明 #

• 两种方案均采用位置标记+批量填充的方式，避免逐行循环，保证处理效率
• data.table方案通过mapply批量生成填充位置，适合组内非缺失值较多的场景
• collapse方案的fby函数针对分组操作做了深度优化，内存占用更低，超大规模数据下表现更优

内容的提问来源于stack exchange，提问作者Maël