Obtendo os melhores valores por grupo

Aqui está um exemplo de quadro de dados:

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

Eu quero que o subconjunto dcontenha as linhas com os 5 principais valores de xpara cada valor de grp.

Usando a base R, minha abordagem seria algo como:

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

Usando dplyr, eu esperava que isso funcionasse:

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

mas ele retorna apenas as 5 primeiras linhas gerais.

Trocar headpor top_nretorna o todo d.

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

Como obtenho o subconjunto correto?

De dplyr 1.0.0 , " slice_min()e slice_max()selecione as linhas com os valores mínimo ou máximo de uma variável, substituindo o confuso 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

Pré- dplyr 1.0.0uso top_n:

De ?top_n, sobre o wtargumento:

A variável a ser usada para o pedido [...] assume a última variável no tbl ".

A última variável em seu conjunto de dados é "grp", que não é a variável que você deseja classificar, e é por isso que sua top_ntentativa "retorna o d inteiro". Portanto, se você deseja classificar por "x" em seu conjunto de dados, é necessário especificar wt = x.

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

Dados:

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

Muito fácil com data.tabletambém ...

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

Ou

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

Ou (deve ser mais rápido para o conjunto de big data porque evita chamar .SDpara cada grupo)

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

Editar: Veja como se dplyrcompara a data.table(se alguém estiver interessado)

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

Adicionando uma data.tablesolução marginalmente mais rápida :

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
)

saída de tempo:

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

Você precisa encerrar headuma chamada para do. No código a seguir, .representa o grupo atual (veja a descrição ...na dopágina de ajuda).

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

Conforme mencionado por akrun, sliceé uma alternativa.

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

Embora eu não tenha perguntado isso, para completar, uma data.tableversão possível é (graças a @Arun pela correção):

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

Minha abordagem na base R seria:

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

E usando dplyr, a abordagem com sliceé provavelmente mais rápida, mas você também pode usar o filterque provavelmente será mais rápido do que usar do(head(., 5)):

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

benchmark 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) ainda retornará várias linhas para cada grupo se a variável de ordenação não for exclusiva dentro de cada grupo. Para selecionar precisamente uma ocorrência para cada grupo, adicione uma variável única a cada linha:

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)

Mais uma data.tablesolução para destacar sua sintaxe concisa:

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