Create countdown and counter before and after a certain value

I want to create a variable that counts down a given number of steps before certain value and counts up a given number steps after the value.

In the example below, I want to have a counters before and after the grp == "b". When grp == "b" the value should be 0, and before and after the counters it should be 100.

To do this I can use case_when function as follows:

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
             )
)

The problem is the variables "before" and "after" are not static! So I should make a function to do it as follows:

b_a<-function(before=15,after=6,others=100,grp_f=letters[2]){

}

How to make a function to do it?

Edit:

I created a function but if there us a better and easy solution please feel free to add an answer:

b_a<-unction(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
} 

This may help:

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))

Try this:

library(dplyr)

test_df %>% 
  group_by(grp) %>% 
  mutate(value = ifelse(grp == letters[index-1] & row_number() <= 5, 100, (before:1)*-1),
         value1 = ifelse(grp == letters[index+1] & row_number() > 3, 100, 1:after)) %>% 
  mutate(value = case_when(grp == letters[index] ~ 0,
                           grp == letters[index-1] ~ value, 
                           grp == letters[index+1] ~ value1), .keep="unused") %>% 
  print(n=50)

  id grp   value
   <int> <chr> <dbl>
 1     1 a       100
 2     2 a       100
 3     3 a       100
 4     4 a       100
 5     5 a       100
 6     6 a        -5
 7     7 a        -4
 8     8 a        -3
 9     9 a        -2
10    10 a        -1
11    11 b         0
12    12 b         0
13    13 b         0
14    14 b         0
15    15 b         0
16    16 b         0
17    17 b         0
18    18 b         0
19    19 b         0
20    20 b         0
21    21 c         1
22    22 c         2
23    23 c         3
24    24 c       100
25    25 c       100
26    26 c       100
27    27 c       100
28    28 c       100
29    29 c       100
30    30 c       100

Set result column to 100. Get indices of the focal group. Create replacement values (countdown - zeros - countup) of given lengths. Replace result at relevant indices.

d = data.frame(id = 1:12, grp = rep(letters[1:3], each = 4))
before = 2
after = 3
g = "b"

d$x = 100
i = which(d$grp == g)
v = c(-before:-1, rep(0, length(i)), seq(after))
d$x = replace(d$x, i[1] - before + 0:(length(v) - 1), v)
d
#    id grp   x
# 1   1   a 100
# 2   2   a 100
# 3   3   a  -2
# 4   4   a  -1
# 5   5   b   0
# 6   6   b   0
# 7   7   b   0
# 8   8   b   0
# 9   9   c   1
# 10 10   c   2
# 11 11   c   3
# 12 12   c 100

Depending on the data, you may want to add a check so that indices of before- and after-sequences are kept within the bounds of the data, e.g. using %in%:

d = data.frame(id = 1:8, grp = rep(letters[1:3], c(1, 5, 2)))
before = 2
after = 3
g = "b"

d$x = 100
i = which(d$grp == g)
v = c(-before:-1, rep(0, length(i)), seq(after))
i2 = i[1] - before + 0:(length(v) - 1)
ok = i2 %in% seq_len(nrow(d)) # <~ check if indices are within range of data 
d$x = replace(d$x, i2[ok], v[ok])
d 
#   id grp  x
# 1  1   a -1 # leading countdown truncated
# 2  2   b  0
# 3  3   b  0
# 4  4   b  0
# 5  5   b  0
# 6  6   b  0
# 7  7   c  1
# 8  8   c  2 # trailing counter truncated