R:通过组X有效识别变量Z的最高N值

Question

我有一个像这样的数据表。

ID <- c(rep("ABC",4), rep("DEF",4), rep("GHI",5))
X  <- c(rep(c(1,2,3,4),3),5)
set.seed(1234)
Z  <- runif(13,min=0, max =1)  
a <- data.table(ID, X, Z)
a
     ID X           Z
 1: ABC 1 0.113703411
 2: ABC 2 0.622299405
 3: ABC 3 0.609274733
 4: ABC 4 0.623379442
 5: DEF 1 0.860915384
 6: DEF 2 0.640310605
 7: DEF 3 0.009495756
 8: DEF 4 0.232550506
 9: GHI 1 0.666083758
10: GHI 2 0.514251141
11: GHI 3 0.693591292
12: GHI 4 0.544974836
13: GHI 5 0.282733584

我想生成一个数据帧，它在每个X子组中具有N个最高的Z值。假设N = 2。我想以这样的数据集结束:

   X  ID         Z
1: 1 DEF 0.8609154
2: 1 GHI 0.6660838
3: 2 DEF 0.6403106
4: 2 ABC 0.6222994
5: 3 GHI 0.6935913
6: 3 ABC 0.6092747
7: 4 ABC 0.6233794
8: 4 GHI 0.5449748
9: 5 GHI 0.2827336

我一直在使用这一行来实现它，但是我发现当数据表很大(即超过1,500,000行或更多)时，它特别慢。

top_n <- 2
a <- a[order(a$X, -a$Z),]
a_2 <- a[, head(.SD, top_n), by=X]
a_2
   X  ID         Z
1: 1 DEF 0.8609154
2: 1 GHI 0.6660838
3: 2 DEF 0.6403106
4: 2 ABC 0.6222994
5: 3 GHI 0.6935913
6: 3 ABC 0.6092747
7: 4 ABC 0.6233794
8: 4 GHI 0.5449748
9: 5 GHI 0.2827336

任何帮助将是最感激的!

谢谢!

Answer 1

这应该比.SD

快

n <- 2
indx <- a[order(-Z), .I[seq_len(n)], by = X]$V1
a[indx]
#      ID  X         Z
#  1: DEF  1 0.8609154
#  2: GHI  1 0.6660838
#  3: GHI  3 0.6935913
#  4: ABC  3 0.6092747
#  5: DEF  2 0.6403106
#  6: ABC  2 0.6222994
#  7: ABC  4 0.6233794
#  8: GHI  4 0.5449748
#  9: GHI  5 0.2827336
# 10:  NA NA        NA

如果你需要一个有序的结果，这也应该是快速的

setorder(a, X, -Z)
indx <- a[, .I[seq_len(n)], by = X]$V1
a[indx]
#      ID  X         Z
#  1: DEF  1 0.8609154
#  2: GHI  1 0.6660838
#  3: DEF  2 0.6403106
#  4: ABC  2 0.6222994
#  5: GHI  3 0.6935913
#  6: ABC  3 0.6092747
#  7: ABC  4 0.6233794
#  8: GHI  4 0.5449748
#  9: GHI  5 0.2827336
# 10:  NA NA        NA