R语言按分组用前值加5向下填充SCORE列NA值的实现问题

数据集 #

原始输入：

HAVE = data.frame(STUDENT = c(1,1,1,2,2,2,3,3),
                  TIME = c(1,2,3,1,2,3,1,2),
                  SCORE = c(7, NA, NA, 5, NA, 19, NA, 2))

期望输出：

WANT = data.frame(STUDENT = c(1,1,1,2,2,2,3,3),
                  TIME = c(1,2,3,1,2,3,1,2),
                  SCORE = c(7, 12, 17, 5, 10, 19, NA, 2))

需求 #

对SCORE列做如下修改：如果值是NA，就用同一STUDENT分组里前一个有效值加5填充，只允许向下填充（不向上补NA）。

原代码的问题 #

用户尝试的代码只能处理单个NA，碰到连续NA就失效：

HAVE %>% group_by(STUDENT) %>% mutate(WANT = ifelse(is.na(SCORE), lag(SCORE) + 5, SCORE))

问题出在lag(SCORE)只取上一行的值，连续NA时，第二个NA对应的lag(SCORE)还是NA，自然算不出正确的填充值。

正确实现方法 #

方法1：结合分组填充与时间差计算 #

先给连续NA段分组，再填充基准值，最后根据TIME间隔计算增量：

library(dplyr)
library(tidyr)

HAVE %>%
  group_by(STUDENT) %>%
  # 标记非NA行，用来划分连续NA的段
  mutate(is_valid = !is.na(SCORE)) %>%
  # 给每个非NA行和后续NA行分配同一个段ID
  mutate(segment_id = cumsum(is_valid)) %>%
  # 向下填充每个段的基准分数和对应TIME
  fill(SCORE, TIME, .direction = "down") %>%
  rename(base_score = SCORE, base_time = TIME) %>%
  # 合并原始数据，计算填充值
  bind_cols(HAVE) %>%
  mutate(WANT = ifelse(is.na(SCORE), base_score + 5*(TIME - base_time), SCORE)) %>%
  select(STUDENT, TIME, SCORE, WANT) %>%
  ungroup()

方法2：用accumulate追踪有效值 #

用purrr的accumulate持续记录最近的有效值和对应TIME，再计算填充值：

library(dplyr)
library(purrr)

HAVE %>%
  group_by(STUDENT) %>%
  mutate(
    # 一直保留最近的非NA分数
    last_valid_score = accumulate(SCORE, ~ ifelse(is.na(.y), .x, .y)),
    # 一直保留最近非NA分数对应的TIME
    last_valid_time = accumulate2(SCORE, TIME, ~ ifelse(is.na(.y), .x, .y)) %>% map_dbl(last),
    # 计算填充值，非NA值直接保留
    WANT = ifelse(is.na(SCORE), last_valid_score + 5*(TIME - last_valid_time), SCORE)
  ) %>%
  select(-last_valid_score, -last_valid_time) %>%
  ungroup()

方法3：简化版（适配初始NA场景） #

针对分组第一个值就是NA的情况（比如学生3的TIME1），保留NA，其他按规则填充：

library(dplyr)
library(tidyr)

HAVE %>%
  group_by(STUDENT) %>%
  mutate(
    # 复制原始列作为基准，后续填充最近的非NA值
    base_score = SCORE,
    base_time = TIME
  ) %>%
  fill(base_score, base_time, .direction = "down") %>%
  # 计算填充值，初始为NA的行保持NA
  mutate(WANT = ifelse(is.na(SCORE) & !is.na(base_score), 
                       base_score + 5*(TIME - base_time), 
                       SCORE)) %>%
  select(STUDENT, TIME, SCORE, WANT) %>%
  ungroup()

以上三种方法都能得到和WANT一致的结果，方法2逻辑更直观，方法1代码更简洁，可根据习惯选择。

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