R语言高效查找每个Type 2行上下最近Type 1行ID的实现方法

以下两种方案均为O(n)时间复杂度的向量化实现，性能远高于手写for循环，可直接用于大规模数据集：

tidyverse 实现 #

library(tidyverse)
# 构造示例数据
df <- tibble(Type=c(1,2,2,1,1,2),ID=c(6,4,3,2,1,5)) 

df_result <- df %>%
  # 生成仅Type=1保留ID的临时列，向下填充得到上方最近的Type1 ID
  mutate(tmp_id = ifelse(Type == 1, ID, NA)) %>%
  fill(tmp_id, .direction = "down") %>%
  rename(IDabove = tmp_id) %>%
  # 同理生成临时列，向上填充得到下方最近的Type1 ID
  mutate(tmp_id = ifelse(Type == 1, ID, NA)) %>%
  fill(tmp_id, .direction = "up") %>%
  rename(IDbelow = tmp_id) %>%
  # Type=1的行对应字段置为NA
  mutate(across(c(IDabove, IDbelow), ~ifelse(Type == 1, NA, .x)))

更高性能的data.table实现（适合超大规模数据集） #

library(data.table)
setDT(df) # 转换为data.table格式

df[, tmp_id := fifelse(Type == 1, ID, NA_integer_)]
df[, IDabove := nafill(tmp_id, type = "locf")] # locf等价于向下填充
df[, tmp_id := fifelse(Type == 1, ID, NA_integer_)]
df[, IDbelow := nafill(tmp_id, type = "nocb")] # nocb等价于向上填充
df[Type == 1, c("IDabove", "IDbelow") := NA_integer_]
df[, tmp_id := NULL] # 删除临时列

两种方案输出结果完全匹配需求：

Type ID IDabove IDbelow
1:    1  6      NA      NA
2:    2  4       6       2
3:    2  3       6       2
4:    1  2      NA      NA
5:    1  1      NA      NA
6:    2  5       1      NA

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