Data.table: select n specific rows before and after other rows that satisfy the condition

Question

Given the following sample data table:

library(data.table)
DT <- fread("grp y exclude
a 1 0
a 2 0
a 3 0
a 4 1
a 5 0
a 7 1
a 8 0
a 9 0
a 10 0
b 1 0
b 2 0
b 3 0
b 4 1
b 5 0
b 6 1
b 7 1
b 8 0
b 9 0
b 10 0
c 5 1
d 1 0")

I want to choose

• by group grp
• all rows having y==5
• and up to two lines before and after each line of 2 in the group.
• but 3. only those lines that have exclude==0.

Assuming each group has a maximum of one row s y==5, this will give the desired result for 1.-3 .:

idx <- -2:2 # 2 rows before match, the matching row itself, and two rows after match
(row_numbers <- DT[,.I[{
                         x <- rep(which(y==5),each=length(idx))+idx 
                         x[x>0 & x<=.N]
                       }], by=grp]$V1)
# [1]  3  4  5  6  7 12 13 14 15 16 20
DT[row_numbers]
#     grp y exclude
#  1:   a 3       0
#  2:   a 4       1
#  3:   a 5       0 # y==5 + two rows before and two rows after
#  4:   a 7       1
#  5:   a 8       0
#  6:   b 3       0
#  7:   b 4       1
#  8:   b 5       0 # y==5 + two rows before and two rows after
#  9:   b 6       1
# 10:   b 7       1
# 11:   c 5       1 # y==5 + nothing, because the group has only 1 element

However, how do I enable 4. to get

#     grp  y exclude
#  1:   a  2       0
#  2:   a  3       0
#  3:   a  5       0
#  4:   a  8       0
#  5:   a  9       0
#  6:   b  2       0
#  7:   b  3       0
#  8:   b  5       0
#  9:   b  8       0
# 10:   b  9       0
# 11:   c  5       1

? I feel like I'm near, but I think I looked at headand whiches too long , so I would be grateful for some fresh ideas.

Answer 1

You are very close. This should do it:

row_numbers <- DT[exclude==0 | y==5, .I[{
    x <- rep(which(y==5), each=length(idx)) + idx 
    x[x>0 & x<=.N]
  }], by=grp]$V1
DT[row_numbers]

Answer 2

:

DT[DT[, rn := .I][exclude==0 | y==5][, rn[abs(.I - .I[y==5]) <= 2], by=grp]$V1]

 #   grp y exclude rn
 #1:   a 2       0  2
 #2:   a 3       0  3
 #3:   a 5       0  5
 #4:   a 8       0  7
 #5:   a 9       0  8
 #6:   b 2       0 11
 #7:   b 3       0 12
 #8:   b 5       0 14
 #9:   b 8       0 17
#10:   b 9       0 18
#11:   c 5       1 20