如何用mutate或其他dplyr动词按规则填充tibble的v1列NA值？

问题背景 #

我们有如下tibble数据框：

library(tidyverse)

set.seed(123) # 固定随机种子保证结果可复现
v1 = c(1 , NA, 3, 5, NA, NA, 2, 12, NA, 5, NA, 0, 1, 2, 6, 8)
alt = rnorm(length(v1), 0, 1)
tb = tibble(v1, alt)

print(tb)

输出结果：

A tibble: 16 × 2
      v1      alt
   <dbl>    <dbl>
 1     1  0.43952
 2    NA  0.76982
 3     3 -0.08362
 4     5  0.07051
 5    NA  0.12929
 6    NA  1.71506
 7     2  0.46092
 8    12 -1.26506
 9    NA -0.68685
10     5  0.83779
11    NA  0.15338
12     0 -1.13814
13     1  1.25382
14     2  0.42646
15     6 -0.29507
16     8  0.89513

填充需求 #

当v1列出现NA时，用**当前行的alt值乘以v1的前一行值（若前一行是已填充后的NA值则沿用该填充值）**进行填充。

已通过for循环实现该逻辑：

tb_for <- tb
for (i in 1:nrow(tb_for)) {
  if( is.na(tb_for[i, 'v1']) ){
    tb_for[i, 'v1'] = tb_for[i-1, 'v1']*tb_for[i, 'alt']
  }
}

print(tb_for)

执行后结果：

A tibble: 16 × 2
        v1      alt
     <dbl>    <dbl>
 1  1.0000  0.43952
 2  0.7698  0.76982
 3  3.0000 -0.08362
 4  5.0000  0.07051
 5  0.6464  0.12929
 6  1.1085  1.71506
 7  2.0000  0.46092
 8 12.0000 -1.26506
 9 -8.2026 -0.68685
10  5.0000  0.83779
11  0.7669  0.15338
12  0.0000 -1.13814
13  1.0000  1.25382
14  2.0000  0.42646
15  6.0000 -0.29507
16  8.0000  0.89513

但for循环在大数据集下效率较低，需要用dplyr或tidyverse工具替代。

高效解决方案：使用purrr::accumulate #

我们可以利用purrr包的accumulate函数实现递推计算，配合dplyr::mutate完成操作，这种向量式操作比循环更高效。

代码实现 #

tb_dplyr <- tb %>%
  mutate(v1 = accumulate(
    .x = 1:n(),
    .f = function(prev_val, idx) {
      if (!is.na(v1[idx])) {
        v1[idx]
      } else {
        prev_val * alt[idx]
      }
    },
    .init = first(v1) # 初始值为第一行的v1
  ) %>% tail(-1)) # 去掉.init添加的初始值，保持长度和原数据一致

print(tb_dplyr)

结果验证 #

运行后得到的tb_dplyr和tb_for结果完全一致，同时处理大数据集时速度远快于for循环。

原理说明 #

• accumulate函数逐个遍历数据行的索引，每次将前一次计算的结果（prev_val）和当前索引（idx）传入自定义函数
• 自定义函数判断当前行的v1是否为NA：非NA则直接使用当前v1值；NA则用前一次结果乘以当前行的alt值
• .init指定初始值为第一行的v1，最后用tail(-1)去掉初始值，保证结果长度与原数据匹配

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