R语言分组过滤数据：代码修正与data.table替代方案问询

数据准备 #

首先创建示例数据框：

id <- c(1,1,1,2,2,3,3,3,3)
x1 <- c("no","yes","yes","no","no","no","no","no","yes")
x2 <- c(9,3,7,4,5,2,5,6,11)
df <- data.frame(id, x1, x2)

需求说明 #

• 按id分组，每个组保留**第一个x1="yes"**的行
• 若组内x1全部为"no"，则保留该组的最后一行

原代码问题分析 #

原代码中if(x1 == "yes")是逐元素判断，生成的逻辑向量长度与分组内行数不匹配，导致筛选逻辑失效，无法得到正确结果。

修正后的dplyr实现 #

library(dplyr)

df %>%
  group_by(id) %>%
  filter(if (any(x1 == "yes")) row_number() == which(x1 == "yes")[1] else row_number() == n()) %>%
  ungroup()

运行后得到期望输出：

# A tibble: 3 × 3
     id x1       x2
  <dbl> <chr> <dbl>
1     1 yes       3
2     2 no        5
3     3 yes      11

data.table替代实现 #

使用data.table的实现方式如下：

library(data.table)

setDT(df)[, .SD[if (any(x1 == "yes")) which(x1 == "yes")[1] else .N], by = id]

输出结果与dplyr版本完全一致。

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