R中高效填充缺失值：基于前序年份线性组合替代循环方案

需要在R中按以下规则填充缺失值：

• 2020年值由2019、2018、2017年值按0.65、0.25、0.10的权重线性组合生成
• 2021年值由2020、2019、2018年值按相同权重组合生成，后续年份以此类推

数据构造代码 #

library(data.table) # CJ函数依赖data.table包
comb1 <- cbind(CJ(letters[1:4], 2000:2019), rnorm(80,2,1)) 
comb2 <-  cbind(CJ(letters[1:4], 2020:2024), data.frame(rep(NA, 20)))
colnames(comb1) <- c("state","year","v")
colnames(comb2) <- c("state","year","v")

comb <- rbind(comb1, comb2)
comb <- comb[order(comb$state, comb$year),]

现有实现的性能问题 #

目前已实现嵌套循环方案，但因数据包含大量state分组，运行速度极慢；后续用向量方式做了部分优化，但仍需更优雅的高效实现方案。

嵌套循环代码 #

library(dplyr)
vlist <- unique(comb$state)
for (i in vlist) {
  for (j in (2002:2024)) {
    value <- 0.65*comb$v[comb$year==j-1 & comb$state==i] + 0.25*comb$v[comb$year==j-2 & comb$state==i] + 0.10*comb$v[comb$year==j-3 & comb$state==i]
    if (is.na(comb$v[comb$year==j & comb$state==i])) {
      comb$v[comb$year==j & comb$state==i] <- value
    }
  }
}

部分优化的向量代码 #

for (j in (2002:2024)) {
  
  comb$v_1 <- ave(comb$v, comb$state, FUN = dplyr::lag)
  comb$v_2 <- ave(comb$v_1, comb$state, FUN = dplyr::lag) 
  comb$v_3 <- ave(comb$v_2, comb$state, FUN = dplyr::lag) 
  
  comb$imp <- 0.65*comb$v_1 + 0.25*comb$v_2 + 0.10*comb$v_3
  comb$v <- ifelse(is.na(comb$v)&comb$year==j, comb$imp, comb$v)
  
  comb$v_1 <- NULL
  comb$v_2 <- NULL
  comb$v_3 <- NULL
  comb$imp <- NULL
}

高效优雅的实现方案 #

方案1：dplyr分组递推填充 #

利用dplyr分组结合accumulate函数实现递推计算，适合习惯tidyverse语法的场景：

library(dplyr)
library(purrr)

comb_filled <- comb %>%
  group_by(state) %>%
  mutate(
    v = accumulate(v, function(last_val, curr_val) {
      if (is.na(curr_val)) {
        # 取当前位置前三个值的加权和
        prev_2 <- lag(v, n=2)[cur_group_rows()][which(v == last_val)]
        prev_3 <- lag(v, n=3)[cur_group_rows()][which(v == last_val)]
        0.65 * last_val + 0.25 * prev_2 + 0.10 * prev_3
      } else {
        curr_val
      }
    })
  ) %>%
  ungroup()

方案2：data.table高效分组计算 #

data.table在大分组数据场景下性能更优，适合数据量较大的情况：

library(data.table)
setDT(comb)

comb[, v := {
  len <- .N
  for (k in 4:len) {
    if (is.na(v[k])) {
      v[k] <- 0.65*v[k-1] + 0.25*v[k-2] + 0.10*v[k-3]
    }
  }
  v
}, by = state]

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