如何用data.table按ID统计break_flag标记行的上下行数

需求实现方案 #

按id分组，将每个break_flag=1的标记行与其前后连续的非标记行（直到下一个标记行或组边界）划分为一个区块，统计每个区块的总行数（标记行自身计1行），并将该数值赋值给区块内所有行。

步骤1：加载并初始化数据

library(data.table)
x <- data.table(id = c(rep("123", 6), rep("234", 4)),
                value = c(1,6,5,10,3,2, 4,6,10,5))
x[, break_flag := ifelse(value == 10, 1, 0)]

步骤2：生成区块并统计行数

# 按id分组，生成基于break_flag的连续分组ID
x[, r := rleid(break_flag), by = id]

# 按id分组，确定包含标记行的分组ID
x[, break_r := unique(r[break_flag == 1]), by = id]

# 构建区块映射：将标记行分组及其前后非标记行分组合并为同一区块
x[, block := {
  r_vals <- unique(r)
  break_rs <- break_r
  block_map <- integer(length(r_vals))
  
  for (br in break_rs) {
    idx <- which(r_vals == br)
    block_map[idx] <- br
    # 合并前一个非标记行分组（如果存在）
    if (idx > 1 && break_flag[r_vals[idx-1] == r][1] == 0) {
      block_map[idx-1] <- br
    }
    # 合并后一个非标记行分组（如果存在）
    if (idx < length(r_vals) && break_flag[r_vals[idx+1] == r][1] == 0) {
      block_map[idx+1] <- br
    }
  }
  
  # 处理无标记行的分组（当前数据无此情况，可选保留）
  if (length(break_rs) == 0) block_map <- rep(1, length(r_vals))
  
  block_map[match(r, r_vals)]
}, by = id]

# 按id和block分组，计算每个区块的总行数
x[, line_count := .N, by = .(id, block)]

最终结果

运行后x的输出如下：

id value break_flag r break_r block line_count
 1: 123     1          0 1       2     2          6
 2: 123     6          0 1       2     2          6
 3: 123     5          0 1       2     2          6
 4: 123    10          1 2       2     2          6
 5: 123     3          0 3       2     2          6
 6: 123     2          0 3       2     2          6
 7: 234     4          0 1       2     2          4
 8: 234     6          0 1       2     2          4
 9: 234    10          1 2       2     2          4
10: 234     5          0 3       2     2          4

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