df 中的 Pandas 样本保持组的平衡答案

【问题标题】：Pandas sample from df keeping balance of groupsdf 中的 Pandas 样本保持组的平衡
【发布时间】：2021-04-05 08:18:57
【问题描述】：

让我们生成一些数据框：

import pandas as pd
categs = ['cat1'] * 600 + ['cat2'] * 300 + ['cat3'] * 100
subcats = ['sub1', 'sub2', 'sub2', 'sub3', 'sub3', 'sub4', 'sub4', 'sub4', 'sub4', 'sub4'] * 100
subcats[0] = 'subX'
vals = range(1000)
df = pd.DataFrame({
   'category': categs,
   'subcategory': subcats,
   'values': vals
})

让我们按类别和子类别查看行数：

print(df.groupby(['category', 'subcategory']).size())

我们得到了

>>>
category  subcategory
cat1      sub1            59
          sub2           120
          sub3           120
          sub4           300
          subX             1
cat2      sub1            30
          sub2            60
          sub3            60
          sub4           150
cat3      sub1            10
          sub2            20
          sub3            20
          sub4            50
dtype: int64

这是一个包含 1000 个元素的数据框。 cat1 有 600 个元素，cat2 有 300 个元素，cat3 有 100 个元素。我想要的是将此数据框从 1000 减少到 60 行，所以
1)每个类别都有相同数量的行（在我们的例子中为 20，等于 60 /（类别数））
2) 保持一个类别中每个子类别的比例
3）如果我们有少量的子类别项目，它仍然留在类别中（在cat1中只有一个'subX'，即使它的比例是cat1的1/600，我们也需要保留它）。

因此，当我们创建新的 df 时，我希望收到这样的信息：

print(newdf.groupby(['category', 'subcategory']).size())


category  subcategory
cat1      sub1            2
          sub2           4
          sub3           4
          sub4           10
          subX             1
cat2      sub1            2
          sub2            4
          sub3            4
          sub4           10
cat3      sub1            2
          sub2            4
          sub3            4
          sub4            10
dtype: int64

在这种情况下，cat1 有 21 个元素，但这没什么大不了的，主要思想是节省了子类别的比例，并且行数在目标数字 20 左右。

【问题讨论】：

您想要每组随机 6%，或者例如每组前 6%？
随机的，不是第一个。而且不会是每组 6%，每个类别需要有相同数量的元素。因此，对于 cat1，它将是 3.33%，对于 cat2，它将是 5.55%，对于 cat3，它将是 20%。

标签： python pandas dataframe pandas-groupby sample

【解决方案1】：

您可以找到每个子类别应保留的行数，并仅保留低于该数的 cumcount 行：

# total (approximate) number of rows to keep
n = 60

# number of rows per category
n_per_cat = n / df['category'].nunique()

# number of rows per subcategory
g_subcat = df.groupby(['category', 'subcategory'])
z = g_subcat['category'].size()
n_per_subcat = np.ceil(z / z.sum(level=0) * n_per_cat)

# output
df_out = (df
          .assign(i=g_subcat.cumcount())
          .merge(n_per_subcat.rename('n').reset_index())
          .query('i < n')
          .drop(columns=['i', 'n']))

# test
df_out.groupby(['category', 'subcategory']).size()

输出：

category  subcategory
cat1      sub1            2
          sub2            4
          sub3            4
          sub4           10
          subX            1
cat2      sub1            2
          sub2            4
          sub3            4
          sub4           10
cat3      sub1            2
          sub2            4
          sub3            4
          sub4           10

附：为了使其随机化，您当然可以在这一切之前对数据帧进行洗牌：

df = df.sample(frac=1)

【讨论】：