如何在熊貓數據框中對數據集進行子集劃分？

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