Pandas根据层次索引选择DataFrameGroupBy中的行

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中的每个用户，我只想保留“probReorder”列中值最大的N行。另外，我希望N依赖于用户\ id。 在我当前的方法中，我有一个dict“lastReordNumber”，它的键值对是（user_id，int），我选择以下行：

^{pr2}$

问题是这太慢了。这个数据帧有大约13M行和200k个唯一的用户的。有没有更快/更好的方法？在

编辑：当给定用户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

您可以将^{}与^{}和^{}一起使用：

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

另一个带有^{}的解决方案：

^{pr2}$