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中的每个user_id，我只想保留"probReorder"列中具有最大值的 N 行。另外，我希望 N 依赖于user_id。 在我目前的方法中，我有一个字典"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)

问题是这真的很慢。数据帧有大约 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