如何在熊猫数据框中对数据集进行子集划分？

Question

我有一个数据框，其中有超过1000万个由大约30列组成的原始数据。

第一列是ID

ID C 
1 1 
1 2 
1 3 
1 2 
1 3 
2 1 
2 5 
2 9 
2 0 
2 1 

我想只提取每个ID的前四行（它们是最新的输入，因为它已经被排序）

我目前正在使用的下面的代码，但不幸的是，它处理大约5％的数据需要大约两个小时的时间，所以处理整个数据可能需要一天左右的时间。

df1 = pd.DataFrame() # an empty dataframe 
for i in df.ID: # df is the dataframe which contains the data 
    df2 = df[df["ID"]== i] 
    df2 = df2[0:4] # take the first four rows 
    df_f = df1.append(df2) 

有没有一种有效的方法可以在短时间内完成同样的事情。

Answer 1

您需要head()方法：

df.groupby("ID").head(4) 

这里是您与运行时间测试原代码的修订版对groupby().head()方法：

def loop(): 
    df1 = pd.DataFrame() # an empty dataframe 
    for i in df.ID.drop_duplicates(): # df is the dataframe which contains the data 
     df2 = df[df["ID"]== i] 
     df2 = df2[0:4] # take the first four rows 
     df1 = pd.concat([df1, df2]) 
    return df1 

%timeit loop() 
# 100 loops, best of 3: 1.99 ms per loop 

%timeit df.groupby("ID").head(4) 
# 1000 loops, best of 3: 485 µs per loop