在R语言中按ID为col_int=1的行添加正负序列的实现方法

针对需求：在数据框中按ID分组，为每组内col_int等于1的位置生成对应序列——该位置之前的行生成递减负向序列，当前行设为1，之后的行生成递增正向序列。以下是可自动化实现的代码，替代示例中的硬编码方式：

示例输入数据 #

library(dplyr)

dfs <- data.frame(
  ID = c(rep("ID1",11), rep("ID2",4), rep("ID3",6)),
  col_int = c(c(0,0,0,0,0,1,0,0,0,0,0), c(1,0,0,0), c(0,0,1,0,1,1))
)

自动化实现代码 #

dfs_processed <- dfs %>%
  group_by(ID) %>%
  mutate(
    # 获取每组中第一个col_int=1的行位置
    target_pos = which(col_int == 1)[1],
    # 生成组内行号索引
    group_row = row_number(),
    # 按规则计算输出序列
    col_out = case_when(
      group_row < target_pos ~ -(target_pos - group_row),
      group_row == target_pos ~ 1,
      group_row > target_pos ~ group_row - target_pos
    )
  ) %>%
  ungroup() %>%
  select(-target_pos, -group_row) # 移除中间变量

代码说明 #

• group_by(ID)：确保每个ID组独立处理。
• target_pos = which(col_int == 1)[1]：锁定每组中第一个标记为1的行作为序列基准点。
• group_row = row_number()：生成组内的相对行号，用于计算序列值。
• case_when分支：
  • 基准点之前的行：用-(target_pos - group_row)生成从-(n-1)到-1的负向序列（n为基准点前的行数）。
  • 基准点行：直接设为1。
  • 基准点之后的行：用group_row - target_pos生成从1开始递增的正向序列。

验证结果 #

运行代码后，生成的col_out与示例中的dfs_out$col_out完全一致：

# 示例输出数据
dfs_out <- dfs %>% mutate(col_out = c(c(seq(-5,-1)),1,c(seq(1,5)), c(1,1,2,3),c(-2,-1,1,1,2,3)))

# 对比验证
all.equal(dfs_processed$col_out, dfs_out$col_out)
# 输出：TRUE

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