如何在R语言数据框中按条件计算下方指定行的col2求和列col3

现有如下R语言数据框：

df <- 
  data.frame(col1 = rep(1:3,2),
             col2 = seq(2,12,2))

数据展示：

col1 col2
1    1    2
2    2    4
3    3    6
4    1    8
5    2   10
6    3   12

需要新增col3列，规则如下：

• 当col1 == 1时，计算该行下方所有col1==2的col2之和；
• 当col1 == 2时，计算该行下方所有col1==3的col2之和；
• 以此类推（实际数据中col1取值范围为1:7）。

期望输出：

col1 col2 col3
1    1    2   14
2    2    4   18
3    3    6   NA
4    1    8   10
5    2   10   12
6    3   12   NA

目前已尝试for循环，希望找到更简便的实现方式。

方法1：dplyr向量化实现 #

利用dplyr的行索引标记，结合purrr::map_dbl完成向量化计算，替代循环：

library(dplyr)
library(purrr)

df <- df %>%
  mutate(row_id = row_number(),
         col3 = map_dbl(row_id, ~{
           target_col1 = col1[.x] + 1
           # 筛选当前行下方、col1等于目标值的col2求和
           sum_val <- sum(col2[row_id > .x & col1 == target_col1], na.rm = FALSE)
           # 无符合条件行时返回NA
           ifelse(sum_val == 0, NA, sum_val)
         })) %>%
  select(-row_id) # 移除临时行索引列

逻辑说明 #

1. 给每行添加唯一row_id，用来定位当前行位置；
2. 对每个行索引，计算当前col1对应的目标值（col1+1）；
3. 筛选出当前行下方、col1匹配目标值的col2数据并求和；
4. 若求和结果为0（无符合条件的行），替换为NA。

方法2：data.table高效实现 #

如果数据量较大，data.table的操作性能更优：

library(data.table)
setDT(df)
df[, row_id := .I] # 添加行索引

df[, col3 := sapply(row_id, function(x) {
  target_col1 = col1[x] + 1
  sum_val <- sum(col2[row_id > x & col1 == target_col1], na.rm = FALSE)
  ifelse(sum_val == 0, NA, sum_val)
})]

df[, row_id := NULL] # 删除临时列

方法3：预分组优化性能 #

针对col1取值连续的特点，预先分组记录各col1对应的行位置，避免重复筛选：

# 先按col1分组，存储每个col1对应的所有行索引
col_positions <- split(seq_len(nrow(df)), df$col1)

df$col3 <- mapply(function(current_row, current_col1) {
  target_col1 = current_col1 + 1
  # 目标col1不存在时直接返回NA
  if (!as.character(target_col1) %in% names(col_positions)) {
    return(NA)
  }
  # 筛选目标col1中在当前行下方的位置
  target_rows <- col_positions[[as.character(target_col1)]]
  target_rows <- target_rows[target_rows > current_row]
  # 有符合条件的行则求和，否则返回NA
  if (length(target_rows) == 0) NA else sum(df$col2[target_rows])
}, seq_len(nrow(df)), df$col1)

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