按组填充数据框首尾非NA观测值间的缺失值

原始数据 #

df <- structure(list(country = c("Slovenia", "Slovenia", "Slovenia", "Slovenia", "Slovenia", "Slovenia", "Hungary", "Hungary", "Hungary", "Hungary", "Hungary", "Hungary"), 
                     year = c(2000, 2001, 2002, 2003, 2004, 2005, 2000, 2001, 2002, 2003, 2004, 2005), 
                     gov_id = c(NA, 1, NA, NA, 2, NA, NA, 12, NA, NA, 13, NA)),
                class = "data.frame", row.names = c("1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"))

需求说明 #

按country分组，用最后一个非NA值向前填充缺失值，但仅填充每组第一个非NA观测值到最后一个非NA观测值之间的行，组开头（如Slovenia的行1、Hungary的行7）和组结尾（如Slovenia的行6、Hungary的行12）的gov_id需保留NA。最终结果如下：

df_new <- structure(list(country = c("Slovenia", "Slovenia", "Slovenia", "Slovenia", "Slovenia", "Slovenia", "Hungary", "Hungary", "Hungary", "Hungary", "Hungary", "Hungary"), 
                         year = c(2000, 2001, 2002, 2003, 2004, 2005, 2000, 2001, 2002, 2003, 2004, 2005), 
                         gov_id = c(NA, 1, 1, 1, 2, NA, NA, 12, 12, 12, 13, NA)),
                    class = "data.frame", row.names = c("1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12"))

问题：直接使用tidyr::fill的缺陷 #

尝试过以下代码：

library(tidyr)
fill(gov_id, direction = "up")

但该方法会错误填充组首尾应保留的NA值，不符合需求。

解决方案 #

方法1：dplyr + tidyr（适合中等规模数据） #

先完成全组填充，再根据原始数据的非NA位置还原首尾区间外的NA：

library(dplyr)
library(tidyr)

df_new <- df %>%
  group_by(country) %>%
  # 保存原始gov_id用于定位非NA行
  mutate(original_gov_id = gov_id) %>%
  # 前向填充所有缺失值
  fill(gov_id, direction = "down") %>%
  # 确定每组第一个和最后一个非NA行的位置
  mutate(
    first_non_na_row = min(which(!is.na(original_gov_id))),
    last_non_na_row = max(which(!is.na(original_gov_id)))
  ) %>%
  # 还原首尾区间外的NA
  mutate(gov_id = case_when(
    row_number() < first_non_na_row ~ NA_real_,
    row_number() > last_non_na_row ~ NA_real_,
    TRUE ~ gov_id
  )) %>%
  # 移除辅助列
  select(-original_gov_id, -first_non_na_row, -last_non_na_row) %>%
  ungroup()

方法2：data.table（适合大规模数据，效率更高） #

利用data.table的分组操作和nafill函数，直接在分组内处理填充与还原：

library(data.table)

setDT(df)

df_new <- df[, {
  original_gov_id <- gov_id
  # 前向填充缺失值
  filled_gov <- nafill(gov_id, type = "locf")
  # 定位首尾非NA行的索引
  first_idx <- min(which(!is.na(original_gov_id)))
  last_idx <- max(which(!is.na(original_gov_id)))
  # 还原首尾区间外的NA
  filled_gov[.I < first_idx | .I > last_idx] <- NA_real_
  # 返回结果
  .(country = country, year = year, gov_id = filled_gov)
}, by = country]

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