如何创建指定分组前后的动态倒计时与计数器函数？

需求说明 #

需要创建一个变量，规则如下：

• 当数据中grp == "b"（或指定分组）时，变量取值为0
• 在该分组前的指定步数（before参数）范围内，进行倒计时（比如分组前第1步为1，第2步为2，直到before步）
• 在该分组后的指定步数（after参数）范围内，进行递增负计数（比如分组后第1步为-1，第2步为-2，直到after步）
• 超出上述范围的行，变量取值为设定的默认值（如100）

静态场景实现 #

已通过dplyr的case_when函数完成固定before和after值的逻辑，代码如下：

library(dplyr)
n<-10;k<-3
test_df<-data.frame(id=1:(k*n),grp=rep(letters[1:k],each=n))
before=5;after=3;index=2
test_df2<-test_df %>%
      mutate(before_after=
          case_when(
             grp == letters[index] ~ 0
             ,dplyr::lag(grp,1) == letters[index] ~ 1
             ,dplyr::lag(grp,2) == letters[index] ~ 2
             ,dplyr::lag(grp,3) == letters[index] ~ 3
             ,dplyr::lead(grp,1) == letters[index] ~ -1
             ,dplyr::lead(grp,2) == letters[index] ~ -2
             ,dplyr::lead(grp,3) == letters[index] ~ -3
             ,dplyr::lead(grp,4) == letters[index] ~ -4
             ,dplyr::lead(grp,5) == letters[index] ~ -5
             ,TRUE~100
             )
)

现有尝试的函数 #

由于before和after参数不固定，需要封装为可动态传参的函数，以下是初步尝试的实现：

b_a<-function(test_df,before=15,after=6,others=100,grp_f=letters[2]){
ltr<-grp_f
test_df2<- test_df %>%
  group_by(grp) %>%
  dplyr::mutate( fr=ifelse(grp==ltr,dplyr::first(id),0 )) %>%
  ungroup() %>%
  dplyr::mutate( fr=max(fr)) %>%
  mutate(fr=-(fr-id)) %>%
  mutate(fr=ifelse(fr<0,fr,0 ))
test_df2<- test_df2 %>%
  group_by(grp) %>%
  dplyr::mutate( lst=ifelse(grp==ltr,dplyr::last(id),0 )) %>%
  ungroup() %>%
  dplyr::mutate( lst=max(lst)) %>%
  mutate(lst=-(lst-id)) %>%
  mutate(lst=ifelse(lst>0,lst,0 ))
test_df2<- test_df2 %>% ungroup() %>%
   mutate(before_after=lst+fr)
test_df2 %>% mutate(ifelse(before_after >= -before & 
    before_after <=after,before_after,others))
test_df2
} 

优化后的简洁实现方案 #

通过定位目标分组的行范围，基于行号计算相对位置，再过滤范围赋值，实现更高效简洁的函数：

library(dplyr)

create_before_after <- function(df, target_grp = "b", before = 5, after = 3, default_val = 100) {
  # 定位目标分组的起始和结束行号
  target_rows <- which(df$grp == target_grp)
  min_target_row <- min(target_rows)
  max_target_row <- max(target_rows)
  
  df %>%
    mutate(
      before_after = case_when(
        # 目标分组内的行赋值0
        between(row_number(), min_target_row, max_target_row) ~ 0,
        # 分组前before步范围内，正计数
        between(row_number(), min_target_row - before, min_target_row - 1) ~ min_target_row - row_number(),
        # 分组后after步范围内，负计数
        between(row_number(), max_target_row + 1, max_target_row + after) ~ (row_number() - max_target_row) * -1,
        # 超出范围赋值默认值
        TRUE ~ default_val
      )
    )
}

# 测试函数
n<-10;k<-3
test_df<-data.frame(id=1:(k*n),grp=rep(letters[1:k],each=n))
result_df <- create_before_after(test_df, target_grp = "b", before = 5, after = 3, default_val = 100)

实现逻辑说明 #

1. 用which定位目标分组的所有行，提取起始和结束行号
2. 通过case_when分四类情况处理：
   • 目标分组内的行直接赋值0
   • 分组前before步的行，用「目标起始行号 - 当前行号」得到正计数值
   • 分组后after步的行，用「当前行号 - 目标结束行号」取反得到负计数值
   • 所有超出范围的行赋值默认值

参考辅助代码 #

用户提供的参考代码片段：

test_df3<-test_df %>%
  group_by(grp) %>%
  mutate(is_group_b=ifelse(grp==letters[index],1,0)) %>%
  ungroup() %>%
  mutate(nr=row_number()) %>%
  mutate(nr=nr*is_group_b) %>%
  mutate(max_nr=max(nr))%>%
  mutate(nr=ifelse(nr==0,NA,nr)) %>%
  mutate(min_nr=min(nr,na.rm = TRUE)) %>%
  mutate(nr=ifelse(is.na(nr),0,nr)) %>%
  mutate(vl=(row_number()-min_nr)) %>%
  mutate(vl=vl*(1-is_group_b)) %>%
  mutate(vl=if_else(vl>0,row_number()-max_nr,vl))

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