如何在pandas数据帧中对数据集进行子集化?

时间:2016-12-06 03:25:20

标签: python python-3.x pandas

我有一个数据框,其中包含超过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) 

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

1 个答案:

答案 0 :(得分:2)

您需要head()方法:

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

enter image description here

以下是原始代码的修订版,其中包含针对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