python - 来自 df 的 Pandas 样本保持群体平衡

Question

让我们生成一些数据帧:

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 左右。

Answer 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)