python - 子集 pandas 数据帧以根据另一个数据帧中的值获取特定行数

Question

我有一个 pandas 数据框，如下所示:

df1

  site_id       date hour  reach                                  maid
0    16002 2023-09-02   21    NaN  33f9fad6-20c5-426c-962f-bc2fbb82aecb
1    16002 2023-09-04   17    NaN  33f9fad6-20c5-426c-962f-bc2fbb82aecb
2    16002 2023-09-04   19    NaN  4a676aeb-6f6f-4622-934b-59b8f149aad7
3    16002 2023-09-04   17    NaN  35363191-c6aa-49fb-beb1-04a98898bed2
4    16002 2023-09-03   22    NaN  a44beb20-a90a-4135-be18-6dda71eeb7c2

我基于上述数据框创建了另一个数据框，该数据框提供每个[site_id,date,hour]组合的记录计数。

df2

      site_id       date hour  count
1666    37226 2023-09-02    8   4586
1676    37226 2023-09-03   16   3586
639     36972 2023-09-03   21   235
640     36972 2023-09-03   22   5431
641     36972 2023-09-03   23   343

我想过滤第一个数据帧并获取第二个数据帧的 count 列中给出的准确记录数。例如，我想从与 site_id 37226、日期 2023-09-02 和小时 8 匹配的第一个数据框中获取 4586 记录。

我尝试在第二个数据帧上使用 forloop，如下所示:

for index,rows in k3.iterrows():
    sid=rows['site_id']
    dt=rows['date']
    hr=rows['hour']
    cnt=rows['count']
    kdf1=dff[(dff['site_id'] == sid) & (dff['date']==dt) & (dff['hour']==hr)]   
    kdf2=kdf1[:cnt]

这可以工作 - 但工作速度非常慢。有没有更快的方法来获取子集。我还附上了两个示例数据框的链接:

Link to df1 and df2

Answer 1

您可以将df2中的count合并到df1，然后使用.groupby减少计数组数:

cols = ["site_id", "date", "hour"]


df1 = df1.merge(df2, on=cols, how="right")
df1 = df1.groupby(cols, group_keys=False).apply(lambda x: x[: x["count"].iloc[0]])
df1.pop("count")

print(df1.head())

打印:

   site_id        date  hour  reach                                  maid
0    37221  2023-09-03    19    NaN  3e769e74-9129-49ba-838d-c36f3a9b3335
1    37221  2023-09-03    19    NaN  71e258d2-5155-4001-9b3c-02a1a1f9c9fb
2    37221  2023-09-03    19    NaN  92eaee88-b41c-4999-b1b8-6be183e5d2cf
3    37221  2023-09-03    19    NaN  c6eb504a-9259-410b-8391-7b06b3e92a41
4    37221  2023-09-03    19    NaN  c36400ff-0790-4844-b58b-2e4cdaafb4d9

注意:对于您的数据，此方法大约需要 0.15 秒，而您的原始版本大约需要 11.2 秒。