如何在data.table中按条件保留/删除分组？有无更优方法？

当然有更高效的实现方式！不用新增临时列再删除，直接利用data.table的分组特性就能搞定，下面给你两种常用的高效写法，再附上性能对比结果。

示例数据 #

先看我们的测试数据集：

set.seed(0)
dt <- data.table(a = rep(1:3, rep(3, 3)), b = sample(1:5, 9, T))
dt
#    a b
# 1: 1 4
# 2: 1 1
# 3: 1 2
# 4: 2 1
# 5: 2 4
# 6: 2 2
# 7: 3 4
# 8: 3 3
# 9: 3 4

我们的需求是保留分组a中包含值2的所有行（也就是删除分组a=3的全部数据）。

方法1：分组内用if(条件) .SD直接返回 #

这是最简洁的写法之一，在分组操作中判断条件，符合要求就返回整个组的.SD（Subset of Data，即当前分组的所有行）：

dt[, if(2 %in% b) .SD, by = a]
#    a b
# 1: 1 4
# 2: 1 1
# 3: 1 2
# 4: 2 1
# 5: 2 4
# 6: 2 2

方法2：用.I获取行索引后子集 #

另一种方式是先分组获取符合条件的行的索引，再用索引直接筛选原表：

dt[dt[, .I[2 %in% b], by = a]$V1]
#    a b
# 1: 1 4
# 2: 1 1
# 3: 1 2
# 4: 2 1
# 5: 2 4
# 6: 2 2

性能对比（基准测试） #

我在2015款Macbook Pro上做了不同数据量级的基准测试，对比上面两种方法的性能：

测试代码：

library(microbenchmark)
library(purrr)

bench <- map(10^(4:7) , ~ {
 dt <- data.table(a = rep(1:.x, rep(10, .x)), b = sample(1:10, 10*.x, T))
 microbenchmark(
   "if+SD" = dt[, if(2 %in% b) .SD, a],
   ".I索引" = dt[dt[, .I[2 %in% b], a]$V1],
   times = 100
 )
})

测试结果 #

• 数据量10^{4个分组（共10}5行）

Unit: milliseconds
    expr      min       lq     mean   median       uq      max neval
   if+SD 13.04827 17.36046 21.15155 19.19119 22.94641  43.0452   100
 .I索引 17.32547 22.92023 27.09775 24.87586 28.39789 108.4760   100

• 数据量10^{5个分组（共10}6行）

Unit: milliseconds
    expr      min       lq     mean   median       uq      max neval
   if+SD 123.9118 143.7802 162.6719 154.4713 173.2986 428.4141   100
 .I索引 158.2975 177.3303 206.3611 193.4460 224.5091 435.3982   100

• 数据量10^{6个分组（共10}7行）

Unit: seconds
    expr      min       lq     mean   median       uq      max neval
   if+SD 1.233100 1.351067 1.448680 1.402827 1.517017 1.852797   100
 .I索引 1.587020 1.704344 1.826468 1.778590 1.947943 2.243176   100

• 数据量10^{7个分组（共10}8行）

Unit: seconds
    expr      min       lq     mean   median       uq      max neval
   if+SD 14.44317 14.65889 14.79806 14.78217 14.91571 15.29134   100
 .I索引 18.04774 18.36764 18.48804 18.45732 18.53073 20.73805   100

结论 #

从测试结果可以看到，if(条件) .SD的写法在所有数据量级下都比.I索引的方法更快，而且代码更简洁，完全不需要临时列，是这类分组筛选需求的最优解之一。对比你原来的“新增列→过滤→删列”的方法，这两种方式不仅代码更清爽，性能优势也会随着数据量增大而愈发明显。

内容的提问来源于stack exchange，提问作者IceCreamToucan