Pandas数据框，根据唯一值的数量选择n个随机行

Question

我正在处理一个文本分类问题，训练效果不错，但我的类别相当不平衡，这妨碍了结果。最大的两个类别比最小的类别大80倍以上，因此有不公平的数量的分类被归为这两个类别。我需要从每个类别中选择n行（任意大）。我的数据集非常大（10m行，1k个唯一类别）。

假设数据框如下：

data = {
    'category':['2','2','2','2','4','4','4','4','4','4','6','6','6'],
    'text':['t1','t2','t3','t4','t5','t6','t7','t8','t9','t10','t11','t12','t13']
}

df = pd.DataFrame(data)

如何从每个分类中选择 n 行随机数据？

我尝试使用 np.random.choice 来选择 n 行随机数据，但我无法找到一种获取索引并进行删除的方法。

当 n = 3 时，理想输出结果应该类似于：

>>> df.head(9)
    category    text
0   2           t3
1   6           t11
2   6           t13
3   4           t6
4   2           t1
5   4           t9
6   4           t8
7   2           t4
8   6           t12

- moses51

Answer 1

你可以使用sample和groupby().head()：

df.sample(frac=1).groupby('category').head(3)

输出：

   category text
4         4   t5
12        6  t13
1         2   t2
8         4   t9
9         4  t10
3         2   t4
10        6  t11
0         2   t1
11        6  t12

- Quang Hoang

---

不错，比我想象的要好得多 :-) - anky