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

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

对于df中的每个用户id，我只想在“probReorder”列中保留具有最大值的N行。另外，我希望N取决于用户的id。 在我当前的方法中，我有一个dict“lastReordNumber”，其键值对是（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)

问题是这真的很慢。数据帧大约有1300万行和20万个唯一用户id。是否有更快/更好的方法

编辑：当给定用户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（）方法的结果。

您可以与and一起使用：

另一个解决方案包括：

---

当两行或更多行等于最大值时会发生什么情况？@BobHaffner问得好。看起来NLAGEST的行为不像我预期的那样，正在复制一些行。我应该发布测试用例的输出吗？我会发布一些包含另一个用户id的额外示例数据，并发布您想要的输出以获得答案。一些注释：在这种情况下，drop_duplicates（）不起任何作用，因为没有重复（用户id、产品id）对。您的第一个解决方案应该与我在编辑中提供的解决方案等效，但它更优雅，可能更高效。您的第二个解决方案在我的机器上无法正常工作，它会产生与我上面提供的相同的“错误”输出。这可能是NLAGEST（）中的错误，我必须查找它。我看到了您的数据，似乎有一些重复。如果没有，那就更好了。如果我的答案有用，别忘了-点击复选标记（

✓），将其从灰显切换为填充。谢谢。正如我所说的（“用户id”、“产品id”）列中没有重复项，如果我错了，请纠正我，这样您取消重复项的呼叫就没有任何作用了。您的两个解决方案相当于我的两个解决方案，但其中一个在我的系统上表现不符合预期。我认为我原来的问题已经解决了，但是我仍然不理解nGuestEnter（）的问题。如果没有重复，只需删除DROPL重复。为什么大型机不起作用？这是一个很难回答的问题。我不知道。也许是虫子。对我来说，它在pandas
0.20.2
中运行良好。你使用最新版本的熊猫吗<代码>打印（pd.show_versions（））

我有版本0.19.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)

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

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