Получение лучших значений по группе

Вот пример фрейма данных:

d <- data.frame(
  x   = runif(90),
  grp = gl(3, 30)
) 

Я хочу, чтобы подмножество dсодержало строки с 5 верхними значениями xдля каждого значения grp.

Используя base-R, мой подход будет примерно таким:

ordered <- d[order(d$x, decreasing = TRUE), ]    
splits <- split(ordered, ordered$grp)
heads <- lapply(splits, head)
do.call(rbind, heads)
##              x grp
## 1.19 0.8879631   1
## 1.4  0.8844818   1
## 1.12 0.8596197   1
## 1.26 0.8481809   1
## 1.18 0.8461516   1
## 1.29 0.8317092   1
## 2.31 0.9751049   2
## 2.34 0.9269764   2
## 2.57 0.8964114   2
## 2.58 0.8896466   2
## 2.45 0.8888834   2
## 2.35 0.8706823   2
## 3.74 0.9884852   3
## 3.73 0.9837653   3
## 3.83 0.9375398   3
## 3.64 0.9229036   3
## 3.69 0.8021373   3
## 3.86 0.7418946   3

Используя dplyr, я ожидал, что это сработает:

d %>%
  arrange_(~ desc(x)) %>%
  group_by_(~ grp) %>%
  head(n = 5)

но он возвращает только верхние 5 строк.

Обмен headна top_nвозвращает все d.

d %>%
  arrange_(~ desc(x)) %>%
  group_by_(~ grp) %>%
  top_n(n = 5)

Как мне получить правильное подмножество?

Начиная с dplyr 1.0.0 , « slice_min()и slice_max()выберите строки с минимальным или максимальным значением переменной, взяв на себя смущение top_n().»

d %>% group_by(grp) %>% slice_max(order_by = x, n = 5)
# # A tibble: 15 x 2
# # Groups:   grp [3]
#     x grp  
# <dbl> <fct>
#  1 0.994 1    
#  2 0.957 1    
#  3 0.955 1    
#  4 0.940 1    
#  5 0.900 1    
#  6 0.963 2    
#  7 0.902 2    
#  8 0.895 2    
#  9 0.858 2    
# 10 0.799 2    
# 11 0.985 3    
# 12 0.893 3    
# 13 0.886 3    
# 14 0.815 3    
# 15 0.812 3

Перед dplyr 1.0.0использованием top_n:

Откуда ?top_n, по поводу wtаргумента:

Переменная, используемая для упорядочивания [...] по умолчанию, соответствует последней переменной в таблице ".

Последняя переменная в вашем наборе данных - это «grp», это не та переменная, которую вы хотите ранжировать, и именно поэтому ваша top_nпопытка «возвращает все d». Таким образом, если вы хотите ранжировать по «x» в вашем наборе данных, вам необходимо указать wt = x.

d %>%
  group_by(grp) %>%
  top_n(n = 5, wt = x)

Данные:

set.seed(123)
d <- data.frame(
  x = runif(90),
  grp = gl(3, 30))

Довольно просто с data.table...

library(data.table)
setorder(setDT(d), -x)[, head(.SD, 5), keyby = grp]

Или

setorder(setDT(d), grp, -x)[, head(.SD, 5), by = grp]

Или (Должно быть быстрее для большого набора данных, потому что не нужно звонить .SDдля каждой группы)

setorder(setDT(d), grp, -x)[, indx := seq_len(.N), by = grp][indx <= 5]

Изменить: вот как можно dplyrсравнить data.table(если кому-то интересно)

set.seed(123)
d <- data.frame(
  x   = runif(1e6),
  grp = sample(1e4, 1e6, TRUE))

library(dplyr)
library(microbenchmark)
library(data.table)
dd <- copy(d)

microbenchmark(
  top_n = {d %>%
             group_by(grp) %>%
             top_n(n = 5, wt = x)},
  dohead = {d %>%
              arrange_(~ desc(x)) %>%
              group_by_(~ grp) %>%
              do(head(., n = 5))},
  slice = {d %>%
             arrange_(~ desc(x)) %>%
             group_by_(~ grp) %>%
             slice(1:5)},
  filter = {d %>% 
              arrange(desc(x)) %>%
              group_by(grp) %>%
              filter(row_number() <= 5L)},
  data.table1 = setorder(setDT(dd), -x)[, head(.SD, 5L), keyby = grp],
  data.table2 = setorder(setDT(dd), grp, -x)[, head(.SD, 5L), grp],
  data.table3 = setorder(setDT(dd), grp, -x)[, indx := seq_len(.N), grp][indx <= 5L],
  times = 10,
  unit = "relative"
)


#        expr        min         lq      mean     median        uq       max neval
#       top_n  24.246401  24.492972 16.300391  24.441351 11.749050  7.644748    10
#      dohead 122.891381 120.329722 77.763843 115.621635 54.996588 34.114738    10
#       slice  27.365711  26.839443 17.714303  26.433924 12.628934  7.899619    10
#      filter  27.755171  27.225461 17.936295  26.363739 12.935709  7.969806    10
# data.table1  13.753046  16.631143 10.775278  16.330942  8.359951  5.077140    10
# data.table2  12.047111  11.944557  7.862302  11.653385  5.509432  3.642733    10
# data.table3   1.000000   1.000000  1.000000   1.000000  1.000000  1.000000    10

Добавление немного более быстрого data.tableрешения:

set.seed(123L)
d <- data.frame(
    x   = runif(1e8),
    grp = sample(1e4, 1e8, TRUE))
setDT(d)
setorder(d, grp, -x)
dd <- copy(d)

library(microbenchmark)
microbenchmark(
    data.table3 = d[, indx := seq_len(.N), grp][indx <= 5L],
    data.table4 = dd[dd[, .I[seq_len(.N) <= 5L], grp]$V1],
    times = 10L
)

вывод времени:

Unit: milliseconds
        expr      min       lq     mean   median        uq      max neval
 data.table3 826.2148 865.6334 950.1380 902.1689 1006.1237 1260.129    10
 data.table4 729.3229 783.7000 859.2084 823.1635  966.8239 1014.397    10

Вам нужно завернуть headв звонок do. В следующем коде .представляет текущую группу (см. Описание ...на doстранице справки).

d %>%
  arrange_(~ desc(x)) %>%
  group_by_(~ grp) %>%
  do(head(., n = 5))

Как уже упоминал Акрун, sliceэто альтернатива.

d %>%
  arrange_(~ desc(x)) %>%
  group_by_(~ grp) %>%
  slice(1:5)

Хотя я не спрашивал об этом, для полноты data.tableкартины возможная версия (спасибо @Arun за исправление):

setDT(d)[order(-x), head(.SD, 5), by = grp]

Мой подход в базе R будет:

ordered <- d[order(d$x, decreasing = TRUE), ]
ordered[ave(d$x, d$grp, FUN = seq_along) <= 5L,]

И при использовании dplyr подход, sliceвероятно, самый быстрый, но вы также можете использовать, filterкоторый, вероятно, будет быстрее, чем использование do(head(., 5)):

d %>% 
  arrange(desc(x)) %>%
  group_by(grp) %>%
  filter(row_number() <= 5L)

тест dplyr

set.seed(123)
d <- data.frame(
  x   = runif(1e6),
  grp = sample(1e4, 1e6, TRUE))

library(microbenchmark)

microbenchmark(
  top_n = {d %>%
             group_by(grp) %>%
             top_n(n = 5, wt = x)},
  dohead = {d %>%
              arrange_(~ desc(x)) %>%
              group_by_(~ grp) %>%
              do(head(., n = 5))},
  slice = {d %>%
             arrange_(~ desc(x)) %>%
             group_by_(~ grp) %>%
             slice(1:5)},
  filter = {d %>% 
              arrange(desc(x)) %>%
              group_by(grp) %>%
              filter(row_number() <= 5L)},
  times = 10,
  unit = "relative"
)

Unit: relative
   expr       min        lq    median        uq       max neval
  top_n  1.042735  1.075366  1.082113  1.085072  1.000846    10
 dohead 18.663825 19.342854 19.511495 19.840377 17.433518    10
  slice  1.000000  1.000000  1.000000  1.000000  1.000000    10
 filter  1.048556  1.044113  1.042184  1.180474  1.053378    10

top_n (n = 1) по-прежнему будет возвращать несколько строк для каждой группы, если переменная порядка не уникальна в каждой группе. Чтобы выбрать ровно одно вхождение для каждой группы, добавьте уникальную переменную в каждую строку:

set.seed(123)
d <- data.frame(
  x   = runif(90),
  grp = gl(3, 30))

d %>%
  mutate(rn = row_number()) %>% 
  group_by(grp) %>%
  top_n(n = 1, wt = rn)

Еще одно data.tableрешение, чтобы выделить его краткий синтаксис:

setDT(d)
d[order(-x), .SD[1:5], grp]