Выберите одну строку в группе с ifelse в data.table

Я группирую таблицу data.table и хочу выбрать из каждой группы первую строку, где x == 1, или, если такая строка не существует, тогда первая строка с любым значением в x

d <- data.table(
           a = c(1,1,1,  2,2,  3,3), 
           x = c(0,1,0,  0,0,  1,1), 
           y = c(1,2,3,  1,2,  1,2)
)

эта попытка

d[, ifelse(any(.SD[,x] == 1),.SD[x == 1][1], .SD[1]), by = a]

возвращает

   a V1
1: 1  1
2: 2  0
3: 3  1

но я ожидал

   a  x  y
1: 1  1  2
2: 2  0  1
3: 3  1  1

Любые идеи, как правильно это сделать?

Ответы

Ответ 1

Мы также можем сделать это с помощью .I, чтобы вернуть индекс строки и использовать его для подмножества строк.

d[d[, .I[which.max(x==1)], by = a]$V1]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1

В текущих версиях data.table подход .I более эффективен по сравнению с .SD для подмножества строк (однако он может измениться в будущем). Это также аналогичный пост

Вот еще один вариант с order (setkey также можно использовать для эффективности) набор данных по 'a' и 'x' после группировки по 'a', а затем получить первую строку с head

d[order(a ,-x), head(.SD, 1) ,by = a]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1

Бенчмарки

Вначале мы думали о бенчмаркинге нa > 1e6, но методы .SD занимают время, поэтому сравнивая строки 3e5 с помощью data.table_1.9.7

set.seed(24)
d1 <- data.table(a = rep(1:1e5, 3), x = sample(0:1, 1e5*3, 
           replace=TRUE), y = rnorm(1e5*3))

system.time(d1[, .SD[which.max(x == 1)], by = a])
#   user  system elapsed 
#  56.21   30.64   86.42 

system.time(d1[, .SD[match(1L, x, nomatch = 1L)], by = a])
# user  system elapsed 
#  55.27   30.07   83.75 

system.time(d1[d1[, .I[which.max(x==1)], by = a]$V1])
#  user  system elapsed 
#   0.19    0.00    0.19 


system.time(d1[order(a ,-x), head(.SD, 1) ,by = a])
# user  system elapsed 
#   0.03    0.00    0.04

Ответ 2

Я думаю, что это хороший вариант использования как для match, так и для аргумента nomatch

d[, .SD[match(1L, x, nomatch = 1L)], by = a]
#    a x y
# 1: 1 1 2
# 2: 2 0 1
# 3: 3 1 1

Это в основном, в случае отсутствия соответствия, возвращает 1, и в результате дает вам первую строку в группе. Если есть множественное совпадение, тогда он вернет первый вариант - по вашему желанию.

Ответ 3

Другая опция (which.max в основном предназначена для того, чтобы делать именно то, что вы хотите):

d[, .SD[which.max(x == 1)], by = a]
#   a x y
#1: 1 1 2
#2: 2 0 1
#3: 3 1 1