按组筛选data.table：满足条件取对应子集，否则取指定子集

需求说明 #

按ID分组筛选数据：

• 若组内存在year=2020的记录，保留该组所有year=2020的行
• 若组内无year=2020的记录，仅保留该组year=0的行

示例数据 #

library(data.table)

# 测试数据
dt <- data.table(ID = c(rep("A",12), rep("B",12), rep("C",6)), 
                 year = c(rep(2020,6), rep(0,6), rep(2019,6), rep(0,12)), 
                 hour = c(rep(1:6,5)), 
                 N = sample(1:20,30, replace = T))

期望输出 #

ID year hour  N
 1:  A 2020    1 12
 2:  A 2020    2 18
 3:  A 2020    3  2
 4:  A 2020    4  7
 5:  A 2020    5 18
 6:  A 2020    6  8
 7:  B    0    1 18
 8:  B    0    2 17
 9:  B    0    3  2
10:  B    0    4  3
11:  B    0    5  9
12:  B    0    6  5
13:  C    0    1 17
14:  C    0    2 19
15:  C    0    3  5
16:  C    0    4 11
17:  C    0    5 20
18:  C    0    6 15

当前实现方式 #

通过两次筛选再合并的方式达成需求：

dt_res1 <- dt[year == 2020]

dt_res2 <- dt[!(ID %in% dt_res1[,ID])][year == 0]

rbindlist(list(dt_res1, dt_res2), use.names = T)

更简洁的分组实现方案 #

利用data.table的分组特性可以一步完成筛选，逻辑更紧凑：

dt[, .SD[if (any(year == 2020)) year == 2020 else year == 0], by = ID]

代码说明 #

• by = ID：按ID字段分组处理数据
• any(year == 2020)：判断当前组是否存在year=2020的记录
• .SD[条件]：根据判断结果筛选当前组内符合要求的行

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