按组选择前n行，每组行数不同

Question

我想选择数据框中每个组的前(2,3,0,4)行。

> f<-data.frame(group=c(1,1,1,2,2,3,4),y=c(1:7))
> 
>   group y
>      1 1
>      1 2
>      1 3
>      2 4
>      2 5
>      3 6
>      4 7

并获得以下数据框：

group y
1 1
1 2
2 4
2 5
4 7

我尝试使用by和head，但是head不接受向量。

谢谢你的帮助。

- Tony

Answer 1

使用更传统的lapply函数：

k <- c(2,3,0,4)
fs <- split(f, f$group)
do.call(rbind,lapply(seq_along(k), function(i) head(fs[[i]], k[i])))

结果是：

  group y
1     1 1
2     1 2
4     2 4
5     2 5
7     4 7

- Aaron left Stack Overflow

---

我差点再次提交了几乎相同的代码...我们的思维方式非常相似。 - Joshua Ulrich

@Joshua：感谢夸奖。我很感激你在这里提供的许多有用答案。顺便说一下，今天早些时候我比你快只是因为你先编辑了问题 - 谢谢你花时间让问题更好。 - Aaron left Stack Overflow

做法很简洁。我有一个大数据框，约有50K个类别在组内。你有任何快速处理的想法吗？干杯 - Tony

@Tony：很难说；我相信split会创建一个新的数据对象，如果是这样的话，而且f非常大，那可能需要一段时间。如果它被证明太慢了，Hadley的plyr包据说可以更快地处理一些类似的事情，你可以尝试一下（就像Sacha的解决方案）。如果你只想要每个组中的几行，也完全有可能通过精心编写的for循环来改进这个问题。另一个可能性是简化split，因为只需要保留y，但这会使记录变得复杂。 - Aaron left Stack Overflow

Answer 2

使用 plyr：

library(plyr)
rows <- c(2,3,0,4)
ddply(f,.(group),function(x)head(x,rows[x[1,1]]))
        group y
    1     1 1
    2     1 2
    3     2 4
    4     2 5
    5     4 7

编辑：

误解问题，现已更新答案。

- Sacha Epskamp

Answer 3

带索引的函数版本。

fun1 <- function(){
  idx <- c(0,which(diff(f$group)!=0))+1
  idx2 <- unlist(lapply(1:length(nf),function(x)  seq.int(from=idx[x],length.out=nf[x])),use.names=F)
  f1 <- f[idx2,]
  return(f1)
}

fun2 <- function(){
  ddply(f,.(group),function(x) head(x,nf[x[1,1]]))
}

测试数据（建议问题作者的大小）

f<-data.frame(group=sample(1:1000,50000,T),y=c(1:50000))
f <- f[order(f$group),]
nf <- rpois(length(unique(f$group)),3) 

system.time(fun1()) system.time(fun2())

在我的系统上，fun1的速度大约比fun2快60倍。

- Wojciech Sobala

---

我使用我的数据集测试了您的代码，确实非常快。对于ddply，我收到以下错误消息“Error in if (n < 0L) max(nrow(x) + n, 0L) else min(n, nrow(x)) : missing value where TRUE/FALSE needed”。 - Tony

你可能在分组变量中有缺失值。你可以将缺失值重新编码为某个值或使用na.omit(f)。 - Wojciech Sobala