Pandas: Select top n groups

Question

假设我有一个像

这样的数据帧

a   b
i1  t1
i1  t2
i2  t3
i2  t1
i2  t3
i3  t2

我想用"a"然后选择2个最大的组。我特别想要生成的行数

a   b
i2  t3
i2  t1
i2  t3
i1  t1
i1  t2

我试着:

df.groupby("a").head(2)   

但是似乎每组选择两行

Answer 1

示例

data = {'a': {0: 'i1', 1: 'i1', 2: 'i2', 3: 'i2', 4: 'i2', 5: 'i3'},
'b': {0: 't1', 1: 't2', 2: 't3', 3: 't1', 4: 't3', 5: 't2'}}
df = pd.DataFrame(data)

代码

---

lst = df['a'].value_counts()[:2].index
out = df[df['a'].isin(lst)]

out

a  b
0   i1  t1
1   i1  t2
2   i2  t3
3   i2  t1
4   i2  t3

---

如果你想按数量排序。使用以下代码

lst = df['a'].value_counts()[:2].index
m = pd.Series(range(0, 2), index=lst)
out = df[df['a'].isin(lst)].sort_values('a', key=lambda x: m[x])

out

a   b
2   i2  t3
3   i2  t1
4   i2  t3
0   i1  t1
1   i1  t2