R语言高效标记事件行及上下指定行数的实现方案问询

示例数据 #

DateTime <- seq(from = as.POSIXct("2021-01-01 00:00:00"), to = as.POSIXct("2021-01-01 17:00:00"), by = "hour")
Binary <- c(NA, 1, rep(NA, 5), 1, rep(NA, 5), 1, rep(NA, 4))
sample <- data.frame(DateTime, Binary)

需求 #

新增Height列，将Binary列为1的行，以及该行上下各1行（可自定义行数）标记为'H'。目标结果如下：

Height <- c(rep('H', 3), rep(NA, 3), rep('H',3), rep(NA, 3), rep('H', 3), rep(NA, 3))
goal <- data.frame(DateTime, Binary, Height)

现有方案的痛点 #

• 用for循环可实现，但百万级数据下速度极慢：

# 创建Height列
sample$Height <- NA
# 循环标记
for (i in 1 : length(sample$Height)){
  if(sample$Binary[i] %in% c(1)){sample$Height [i] <- "H"}
  if(sample$Binary[i] %in% c(1)){sample$Height [i+1] <- "H"}
  if(sample$Binary[i] %in% c(1)){sample$Height [i-1] <- "H"}
} 

• dplyr仅能标记Binary=1的行，无法覆盖上下行；fill()会填充所有NA，不符合需求：

sample <- sample %>%
  mutate(Height = ifelse(Binary==1,'H', NA))

高效解决方案 #

方案1：dplyr + 索引扩展 #

通过行索引批量处理，避免循环：

library(dplyr)

# 定义上下扩展行数
n <- 1
# 获取Binary=1的行索引
target_idx <- which(sample$Binary == 1)
# 扩展索引范围，过滤超出数据框边界的行
expanded_idx <- unique(unlist(lapply(target_idx, function(x) seq(max(1, x-n), min(nrow(sample), x+n)))))
# 批量标记
sample <- sample %>%
  mutate(Height = ifelse(row_number() %in% expanded_idx, 'H', NA))

方案2：zoo包滑动窗口 #

利用滑动窗口检查每行周围是否存在事件：

library(dplyr)
library(zoo)

n <- 1
# 窗口大小：当前行+上下n行
window_size <- 2*n + 1

sample <- sample %>%
  mutate(Height = ifelse(
    rollapply(Binary, window_size, function(x) any(x == 1, na.rm = TRUE), 
              fill = FALSE, align = "center"), 
    'H', NA
  ))

方案3：data.table（大数据最优解） #

data.table的向量化操作在百万级数据下效率最高：

library(data.table)

setDT(sample)
n <- 1
# 获取目标行索引
target_idx <- sample[Binary == 1, .I]
# 扩展索引
expanded_idx <- unique(unlist(lapply(target_idx, function(x) seq(max(1, x-n), min(nrow(sample), x+n)))))
# 批量赋值
sample[expanded_idx, Height := 'H']

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