python - 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