在data.table中使用ifelse为每个组选择一行

Question

我正在对data.table进行分组,并希望从每个组中选择x == 1的第一行,或者如果这样的行不存在,那么第一行中的任何值都是x

d <- data.table(
           a = c(1,1,1,  2,2,  3,3), 
           x = c(0,1,0,  0,0,  1,1), 
           y = c(1,2,3,  1,2,  1,2)
)

这种尝试

d[, ifelse(any(.SD[,x] == 1),.SD[x == 1][1], .SD[1]), by = a]

回报

   a V1
1: 1  1
2: 2  0
3: 3  1

但我期待

   a  x  y
1: 1  1  2
2: 2  0  1
3: 3  1  1

任何想法如何正确吗？

Answer 1

我认为这对两者都是一个很好的用例,match 而且它的nomatch论点

d[, .SD[match(1L, x, nomatch = 1L)], by = a]
#    a x y
# 1: 1 1 2
# 2: 2 0 1
# 3: 3 1 1

基本上,如果不匹配,则返回1,结果为您提供组中的第一行.如果存在多重匹配,那么它将根据您的愿望返回第一个匹配

Answer 2

另一种选择(which.max基本上设计为完全按照你的意愿):

d[, .SD[which.max(x == 1)], by = a]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1

Answer 3

我们也可以这样做.I来返回行索引并使用它来对行进行子集化.

d[d[, .I[which.max(x==1)], by = a]$V1]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1

在当前版本中data.table,.I与.SD子集化行相比,方法更有效(但是,它可能在将来发生变化).这也是一个类似的帖子

这是另一个选项order(setkey也可以用于 - 效率)数据集'a'和'x'分组后'a',然后得到第一行head

d[order(a ,-x), head(.SD, 1) ,by = a]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1

基准

最初,我们考虑在> 1e6上进行基准测试,但这些.SD方法需要时间,所以比较使用的3e5行data.table_1.9.7

set.seed(24)
d1 <- data.table(a = rep(1:1e5, 3), x = sample(0:1, 1e5*3, 
           replace=TRUE), y = rnorm(1e5*3))

system.time(d1[, .SD[which.max(x == 1)], by = a])
#   user  system elapsed 
#  56.21   30.64   86.42 

system.time(d1[, .SD[match(1L, x, nomatch = 1L)], by = a])
# user  system elapsed 
#  55.27   30.07   83.75 

system.time(d1[d1[, .I[which.max(x==1)], by = a]$V1])
#  user  system elapsed 
#   0.19    0.00    0.19 


system.time(d1[order(a ,-x), head(.SD, 1) ,by = a])
# user  system elapsed 
#   0.03    0.00    0.04