r - data.table:在满足条件的其他行之前和之后选择n个特定行

Question

给定以下示例数据表:

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

我要选择

按组grp

所有具有y==5的行

以及分组中2之内的每一行之前和之后的最多两行。

，但仅3.具有exclude==0的行。

假设每个组最多包含一行 y==5，这将为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

但是，我该如何合并4.以便获得

#     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

？感觉我已经接近了，但是现在我猜想我对 head和 which s的查找时间太长了，所以我要感谢一些新的想法。

Answer 1

你很亲密这应该做到这一点:

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]