R语言：用0或上方值n重复n次填充列中的NA值

需求说明 #

给定如下数据框：

df <- data.frame(x = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12),
                 y = c(NA, 2, NA, NA, NA, 3, NA, NA, NA, 1, NA, NA))

需要将其转换为以下结果：

data.frame(x = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12),
           y = c(0, 2, 2, 0, 0, 3, 3, 3, 0, 1, 0, 0))

#>     x y
#> 1   1 0
#> 2   2 2
#> 3   3 2
#> 4   4 0
#> 5   5 0
#> 6   6 3
#> 7   7 3
#> 8   8 3
#> 9   9 0
#> 10 10 1
#> 11 11 0
#> 12 12 0

核心规则：

1. 所有NA替换为0
2. 若某位置的值>=2，则向后填充该值减1次（比如值为2则填充后面1行，值为3则填充后面2行）
3. 值为1时不做填充操作

已通过while循环实现需求，现寻求更符合R语言风格的解决方案，同时支持管道操作与多列处理。

用户原循环实现代码：

df[is.na(df)] <- 0 # replace all NA with 0

i = 1
while (i < nrow(df)){
  
  if (df$y[i] < 2){ # do nothing if y = 1
    i = i+1
  } else {
    df$y[(i+1):(i+df$y[i]-1)] <- df$y[i]
    i = i+df$y[i]
  }
}

解决方案 #

1. 单列处理（管道风格） #

利用dplyr结合向量操作，实现无循环的管道式处理：

library(dplyr)

df_processed <- df %>%
  # 第一步：替换NA为0
  mutate(y = replace(y, is.na(y), 0)) %>%
  # 第二步：批量处理填充逻辑
  {
    # 获取需要填充的起始位置和对应填充长度
    fill_positions <- which(.$y >= 2)
    fill_lengths <- .$y[fill_positions] - 1
    
    # 生成需要替换的行索引，确保不超出数据框范围
    replace_indices <- unlist(mapply(function(pos, len) pos + 1:len, fill_positions, fill_lengths))
    replace_indices <- replace_indices[replace_indices <= nrow(.)]
    
    # 执行填充
    .$y[replace_indices] <- .$y[rep(fill_positions, fill_lengths)]
    .
  }

2. 多列处理（管道风格） #

通过自定义处理函数+dplyr::across实现多列批量处理，以新增z列为例：

library(dplyr)

# 定义通用填充函数
fill_custom <- function(vec) {
  # 替换NA为0
  vec <- replace(vec, is.na(vec), 0)
  
  # 无需要填充的情况直接返回
  fill_positions <- which(vec >= 2)
  if (length(fill_positions) == 0) return(vec)
  
  # 生成填充索引并执行替换
  fill_lengths <- vec[fill_positions] - 1
  replace_indices <- unlist(mapply(function(pos, len) pos + 1:len, fill_positions, fill_lengths))
  replace_indices <- replace_indices[replace_indices <= length(vec)]
  vec[replace_indices] <- vec[rep(fill_positions, fill_lengths)]
  
  return(vec)
}

# 带多列的原始数据框
df_multi <- data.frame(
  x = c(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12),
  y = c(NA, 2, NA, NA, NA, 3, NA, NA, NA, 1, NA, NA),
  z = c(1, NA, NA, NA, 4, NA, NA, NA, NA, 2, NA, NA)
)

# 管道式多列处理
df_multi_processed <- df_multi %>%
  mutate(across(c(y, z), fill_custom))

执行后df_multi_processed的z列结果为：c(1, 0, 0, 0, 4, 4, 4, 4, 0, 2, 2, 0)，符合需求。

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