Question

我有以下数据框（称为 df ）：

   user_id  product_id  probReorder
0        1         196          1.0
1        1       10258          0.9
2        1       10326          0.1
3        1       12427          1.0
4        1       13032          0.3
...

对于 df 中的每个 user_id ，我只想保留＆＃34; probReorder＆＃34;中的最大值的N行。柱。另外，我希望N依赖 user_id 。在我目前的方法中，我有一个字典＆＃34; lastReordNumber＆＃34;其键值对是（user_id，int），我按如下方式选择行：

predictions = []
for usr,data in df.groupby(by="user_id"):
    data = data.nlargest(lastReordNumber[usr], "probReorder")
    predictions.append(data)
df = pd.concat(predictions)

问题是这真的很慢。数据框有大约13M行和200k唯一 user_id 。是否有更快/更好的方法？

编辑：当给定 user_id 的 probReorder 列中存在重复值时，上一代码会产生意外输出。例如：

lastReordNumber = {1:2, 2:3}
df = pd.DataFrame({"user_id":[1,1,1,2,2,2,2],"probReorder":[0.9,0.6,0.9,0.1,1,0.5,0.4],\
    "product_id":[1,2,3,4,5,6,7]})

我得到了输出：

   probReorder  product_id  user_id
0          0.9           1        1
1          0.9           3        1
2          0.9           1        1
3          0.9           3        1
4          1.0           5        2
5          0.5           6        2
6          0.4           7        2

对于user_id = 2是我期望的，但对于user_id = 1，有重复的行。我的预期输出是：

   probReorder  product_id  user_id
0          0.9           1        1
1          0.9           3        1
2          1.0           5        2
3          0.5           6        2
4          0.4           7        2

这可以通过使用更简单的代码来获得

predictions = []
for usr,data in df.groupby(by="user_id"):
    predictions.append(data.sort_values('probReorder', ascending=False).head(lastReordNumber[usr]))
predictions = pd.concat(predictions, ignore_index=True)

其中每列完全排序然后被截断。这也是相当有效的。但我还没有理解如何解释nlargest（）方法的结果。

Answer 1

您可以将sort_values与groupby和head：

一起使用

df1 = df.sort_values('probReorder', ascending=False)
        .groupby('user_id', group_keys=False)
        .apply(lambda x: x.head([x.name]))
print (df1)
   probReorder  product_id  user_id
0          0.9           1        1
2          0.9           3        1
4          1.0           5        2
5          0.5           6        2
6          0.4           7        2

nlargest的另一个解决方案：

df1 = df.groupby('user_id', group_keys=False)
        .apply(lambda x: x.nlargest(lastReordNumber[x.name], 'probReorder'))
print (df1)
   probReorder  product_id  user_id
0          0.9           1        1
2          0.9           3        1
4          1.0           5        2
5          0.5           6        2
6          0.4           7        2

Pandas基于分层索引

1 个答案: