基于辅助列对分组缺失值进行前后向填充的dplyr实现方法

问题描述 #

假设我有如下数据：

input = tibble::tibble(
  group = c(rep("A", 5), rep("B", 5), rep("C", 5)),
  value = c(10, 15, 17, NA, NA, NA, NA, 12, 16, 13, 12, NA, 15, NA, 19),
  gr = c(0.1, 0.05, 0.03, 0.02, 0.05, 0.04, 0.02, 0.6, 0.03, 0.4, 0.01, 0.09, 0.05, -0.03, 0.04)
)

数据结构如下：

> input
# A tibble: 15 x 3
   group value    gr
   <chr> <dbl> <dbl>
 1 A        10  0.1 
 2 A        15  0.05
 3 A        17  0.03
 4 A        NA  0.02
 5 A        NA  0.05
 6 B        NA  0.04
 7 B        NA  0.02
 8 B        12  0.6 
 9 B        16  0.03
10 B        13  0.4 
11 C        12  0.01
12 C        NA  0.09
13 C        15  0.05
14 C        NA -0.03
15 C        19  0.04

我需要使用辅助变量gr对每个group分组的value列缺失值进行填充，不同分组的填充规则不同：

• groupA采用向前填充，计算公式为value_filled = lag(value) * (1 + gr)
• groupB采用向后填充，计算公式为value_filled = lag(value) / (1 + gr)
• groupC的缺失值位于有效取值之间，采用向前填充即可。

期望得到的输出结果如下：

desired_output = tibble::tibble(
  group = c(rep("A", 5), rep("B", 5), rep("C", 5)),
  value = c(10, 15, 17, NA, NA, NA, NA, 12, 16, 13, 12, NA, 15, NA, 19),
  gr = c(0.1, 0.05, 0.03, 0.02, 0.05, 0.04, 0.02, 0.6, 0.03, 0.4, 0.01, 0.09, 0.05, -0.03, 0.04),
  value_filled = c(10, 15, 17, 17.3, 18.2, 7.3, 7.5, 12, 16, 13, 12, 13, 15, 14.5, 19)
)

> desired_output
# A tibble: 15 x 4
   group value    gr value_filled
   <chr> <dbl> <dbl>        <dbl>
 1 A        10  0.1          10  
 2 A        15  0.05         15  
 3 A        17  0.03         17  
 4 A        NA  0.02         17.3
 5 A        NA  0.05         18.2
 6 B        NA  0.04         7.3
 7 B        NA  0.02         7.5
 8 B        12  0.6          12  
 9 B        16  0.03         16  
10 B        13  0.4          13  
11 C        12  0.01         12  
12 C        NA  0.09         13  
13 C        15  0.05         15  
14 C        NA -0.03         14.5
15 C        19  0.04         19 

请提供符合dplyr语法风格的实现代码。

解决代码 #

实现思路 #

• 三类分组的填充逻辑均为迭代依赖上一步的填充结果，因此选用purrr::accumulate2()实现滚动迭代计算
• 分组A：缺失值在序列末尾，正向迭代即可，非缺失值保留原值，缺失值用上一步填充值乘以(1+当期gr)
• 分组B：缺失值在序列开头，先将序列倒序后迭代，非缺失值保留原值，缺失值用上一步填充值除以(1+当期gr)，迭代完成后恢复原顺序
• 分组C：缺失值在序列中部，正向迭代规则与分组A一致

代码实现 #

# 加载依赖包
library(tidyverse)

result <- input %>%
  # 按分组处理
  group_by(group) %>%
  mutate(value_filled = case_when(
    # 处理A组
    group == "A" ~ accumulate2(value, gr, ~ ifelse(is.na(..2), ..1 * (1 + ..3), ..2)) %>% unlist(),
    # 处理B组
    group == "B" ~ {
      rev_val <- rev(value)
      rev_gr <- rev(gr)
      accumulate2(rev_val, rev_gr, ~ ifelse(is.na(..2), ..1 / (1 + ..3), ..2)) %>% 
        unlist() %>% 
        rev()
    },
    # 处理C组
    group == "C" ~ accumulate2(value, gr, ~ ifelse(is.na(..2), ..1 * (1 + ..3), ..2)) %>% unlist()
  )) %>%
  ungroup() %>%
  # 按示例保留1位小数，可根据需求调整精度
  mutate(value_filled = round(value_filled, 1))

运行上述代码得到的result与期望输出完全一致。

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