如何按组筛选/缩减数据集？（data.table实现）

我有一个包含多列和分组变量grp的数据集，希望按分组变量缩减数据集，每个分组最多保留max_n行，同时保留其他列的分布（即保留筛选后a、b列的最小值和最大值），因此先使用了setorderv函数排序。

示例数据集代码：

library(data.table)

set.seed(22)
n=20
max_n = 6
dt <- data.table("grp"=sample(c("a", "b", "c"), n, replace=T),
                 "a"=sample(1:10, n, replace=T),
                 "b"=sample(1:20, n, replace=T),
                 "id"=1:n)
setorderv(dt, c("grp", "a", "b"))
dt

当前的临时解决方案不够优雅，不符合data.table风格：

dt_new <- data.table()
for (gr in unique(dt[["grp"]])) {
  tmp <- dt[grp == gr, ]
  n_tmp <- nrow(tmp)
  if (n_tmp > max_n) {
    tmp <- tmp[as.integer(seq(1, n_tmp, length.out=max_n)),]
  }
  dt_new <- rbindlist(list(dt_new, tmp))
}

需要更优雅的data.table实现方式。

直接利用data.table的分组操作结合内置变量与函数，一行代码即可完成，完全贴合data.table的向量化风格：

dt_new <- dt[, .SD[as.integer(seq.int(1, .N, length.out = min(.N, max_n)))], by = grp]

关键细节说明 #

• .N：分组语境下代表当前分组的总行数
• min(.N, max_n)：自动适配分组行数，小于等于max_n时保留全部行，超过时按指定数量筛选
• seq.int(1, .N, length.out = ...)：生成等间隔的行索引，确保覆盖分组首尾行（满足保留a、b列最小/最大值的需求）
• .SD：指代当前分组的子数据集，通过索引直接筛选目标行

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