在data.table中向前填充前置NA且保留后置NA的方法

问题描述 #

在data.table中需要实现：将每列里第一个非NA值之后的NA填充为最近的前置非NA值，但保留第一个非NA值之前的NA，以及非NA值后续出现的NA（比如某列新出现非NA值后，其后续的NA保持不变）。zoo::na.locf()会把所有后续NA都填充，无法满足需求。

数据为区域层级映射表，region_level1是顶层父项，后续列是细分区域：部分父项有多个子项（如Cerebral cortex分为Cerebral cortex 1和Cerebral cortex 2），而Temporal cortex无细分，因此后续列存在NA。

输入数据 #

input_data <- data.frame(
  region_level1 = c("Brain", NA, NA, NA, NA, NA),
  region_level2 = c(NA, "Grey Matter", NA, NA, NA, NA),
  region_level3 = c(NA, NA, "Cerebral Cortex", NA, NA, "Temporal Cortex"),
  region_level4 = c(NA, NA, NA, "Cerebral cortex 1", "Cerebral cortex 2", NA),
  stringsAsFactors = FALSE
)

输入数据预览：

region_level1 region_level2     region_level3       region_level4
1        Brain          <NA>              <NA>                <NA>
2         <NA>  Grey Matter              <NA>                <NA>
3         <NA>         <NA> Cerebral Cortex                <NA>
4         <NA>         <NA>              <NA> Cerebral cortex 1
5         <NA>         <NA>              <NA> Cerebral cortex 2
6         <NA>         <NA> Temporal Cortex                <NA>

期望输出 #

desired_output <- data.frame(
  region_level1 = c("Brain", "Brain", "Brain", "Brain", "Brain", "Brain"),
  region_level2 = c(NA, "Grey Matter", "Grey Matter", "Grey Matter", "Grey Matter", "Grey Matter"),
  region_level3 = c(NA, NA, "Cerebral Cortex", "Cerebral Cortex", "Cerebral Cortex", "Temporal Cortex"),
  region_level4 = c(NA, NA, NA, "Cerebral cortex 1", "Cerebral cortex 2", NA),
  stringsAsFactors = FALSE
)

期望结果预览：

region_level1 region_level2     region_level3       region_level4
1        Brain          <NA>              <NA>                <NA>
2        Brain  Grey Matter              <NA>                <NA>
3        Brain  Grey Matter Cerebral Cortex                <NA>
4        Brain  Grey Matter Cerebral Cortex Cerebral cortex 1
5        Brain  Grey Matter Cerebral Cortex Cerebral cortex 2
6        Brain  Grey Matter Temporal Cortex                <NA>

na.locf()的局限性 #

使用zoo::na.locf()会填充所有后续NA，比如第6行的region_level4被错误填充为Cerebral cortex 2，不符合需求：

library(zoo)
converted_data <- na.locf(input_data)

结果预览：

region_level1 region_level2     region_level3       region_level4
1        Brain          <NA>              <NA>                <NA>
2        Brain  Grey Matter              <NA>                <NA>
3        Brain  Grey Matter Cerebral Cortex                <NA>
4        Brain  Grey Matter Cerebral Cortex Cerebral cortex 1
5        Brain  Grey Matter Cerebral Cortex Cerebral cortex 2
6        Brain  Grey Matter Temporal Cortex Cerebral cortex 2

解决方案 #

方法1：data.table原生遍历实现 #

通过标记每列的非NA位置，精准填充对应范围的NA：

library(data.table)

# 转换为data.table格式
dt <- as.data.table(input_data)

# 遍历每一列处理
for (col in names(dt)) {
  # 获取当前列所有非NA的行索引
  non_na_pos <- which(!is.na(dt[[col]]))
  if (length(non_na_pos) > 0) {
    # 为每个非NA值确定填充范围
    for (i in seq_along(non_na_pos)) {
      start_row <- non_na_pos[i]
      # 下一个非NA的前一行，若无则取最后一行
      end_row <- if (i < length(non_na_pos)) non_na_pos[i+1] - 1 else nrow(dt)
      # 填充start_row到end_row的NA为当前非NA值
      dt[start_row:end_row, (col) := dt[[col]][start_row]]
    }
  }
}

# 查看结果
dt

方法2：na.locf结合原始NA还原 #

先全量填充，再还原第一个非NA值之前的NA，更简洁：

library(zoo)
library(data.table)

dt <- as.data.table(input_data)

# 保存原始数据的NA位置标记
original_na_flags <- dt[, lapply(.SD, is.na)]

# 用na.locf填充所有NA
filled_dt <- dt[, lapply(.SD, na.locf)]

# 还原每列第一个非NA值之前的原始NA
for (col in names(dt)) {
  first_non_na_idx <- which(!is.na(dt[[col]]))[1]
  if (!is.na(first_non_na_idx)) {
    filled_dt[1:(first_non_na_idx - 1), (col) := NA]
  }
}

# 查看结果
filled_dt

方法3：Base R实现 #

无需额外包，用分组填充实现需求：

# 定义单列填充函数
fill_target_na <- function(col_vector) {
  # 生成分组：每个非NA值开启一个新组，前置NA为组0
  groups <- cumsum(!is.na(col_vector))
  # 按组填充：组内统一为第一个非NA值，组0保留NA
  ave(col_vector, groups, FUN = function(g) {
    if (length(g) == 0 || is.na(g[1])) NA else g[1]
  })
}

# 对所有列应用函数
filled_data <- as.data.frame(lapply(input_data, fill_target_na))

# 查看结果
filled_data

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