Pandas 来自 df 的样本保持组平衡

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'，即使它的比例是1/我们也需要保留它cat1 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

P.S。为了使其随机，当然，您可以在所有这些之前随机播放数据帧：

df = df.sample(frac=1)

Pandas 来自 df 的样本保持组平衡

Pandas sample from df keeping balance of groups

python

pandas

dataframe

pandas-groupby

sample