如何在 Pandas 数据框中对数据集进行子集筛选？

Question

我有一个数据框，它由约30列组成，超过1000万行。

第一列是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) 

有没有一种更有效的方式来在更短的时间内完成同样的事情？

- Natheer Alabsi

---

是否保证每个ID至少有四个实例？ - kiliantics

是的，大多数ID都有超过10个实例，我只想获取最近四个月的实例，并且每个ID的实例已按降序排序。 - Natheer Alabsi

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

- Psidom

---

我使用了你的代码：df.groupby("ID").head(4)它解决了我的问题，而且没有使用循环。非常感谢。 - Natheer Alabsi