Pandas分组只选取前N个组

Question

我有一些DataFrame，想按ID分组，例如：

import pandas as pd
df = pd.DataFrame({'item_id': ['a', 'a', 'b', 'b', 'b', 'c', 'd'], 'user_id': [1,2,1,1,3,1,5]})
print df

生成：

  item_id  user_id
0       a        1
1       a        2
2       b        1
3       b        1
4       b        3
5       c        1
6       d        5

[7 rows x 2 columns]

我可以轻松地按id分组：

grouped = df.groupby("item_id")

但是我如何只返回前N个分组对象？例如，我只想要前3个唯一的item_id。

- Christian Sauer

---

先筛选数据框(df)再进行操作，会不会更容易呢？可以这样写：df[df['item_id'].isin(df['item_id'].unique()[:3])].groupby('item_id')。该代码会先根据 'item_id' 列中出现的前三个唯一值，对数据框进行筛选，并按照 'item_id' 列进行分组。 - EdChum

@firelynx 我同意。我之前没有找到那个结果。 - Christian Sauer

迭代前三个组。其中g是groupby的一个实例，代码如下：for n,(k,gg) in enumerate(list(g)[:3])。 - BSalita

Answer 1

这里有一种使用 list(grouped) 的方法。

result = [g[1] for g in list(grouped)[:3]]

# 1st
result[0]

  item_id  user_id
0       a        1
1       a        2

# 2nd
result[1]

  item_id  user_id
2       b        1
3       b        1
4       b        3

- Jianxun Li

---

谢谢，那是个好主意。由于一些限制，最后我使用了一个随机查询。 - Christian Sauer

这会实例化列表中的所有组，而你只需要前3个，因此对于大型数据帧来说非常低效。 - Denziloe

Answer 2

一种方法是使用Counter从列表中获取前三个唯一项，根据这些项过滤您的DataFrame，然后在此过滤后的DataFrame上执行groupby操作。

from collections import Counter

c = Counter(df.item_id)
most_common = [item for item, _ in c.most_common(3)]

>>> df[df.item_id.isin(most_common)].groupby('item_id').sum()
         user_id
item_id         
a              3
b              5
c              1

- Alexander

---

谢谢，那是个好主意。由于一些限制，最后我使用了一个随机查询。 - Christian Sauer