Python pandas:有条件地从数据帧中选择一个统一的样本

Question

假设我有一个数据帧

category1  category2   other_col   another_col ....
a          1
a          2
a          2        
a          3
a          3
a          1
b          10
b          10
b          10
b          11
b          11
b          11

我想从我的数据帧中获取样本,以便category1统一次数.我假设每种类型的数量相同category1.我知道这可以通过使用pandas来完成pandas.sample().但是,我还想确保我选择的样本也category2同样具有代表性.所以,例如,如果我的样本大小为5,我会想要一些如下:

a  1
a  2
b  10
b  11
b  10

我不想要这样的东西:

a 1
a 1
b 10
b 10
b 10

虽然这是一个有效的随机样本n=4,它不会因为我想改变尽可能的类型,满足我的要求category2.

请注意,在第一个示例中,因为a仅采样了两次,3因此未表示category2.这没关系.目标是尽可能均匀地表示样本数据.

如果它有助于提供更清晰的例子,我们可以有事情的类别fruit,vegetables,meat,grains,junk.在10的样本大小中,我希望尽可能多地代表每个类别.理想情况下,每个2个.然后,属于所选类别的那两个所选行中的每一个将具有也被尽可能均匀地表示的子类别.因此,例如,水果可以具有red_fruits,yellow_fruits等的子类别.对于从10中选择的2个水果类别,red_fruits和yellow_fruits都将在样本中表示.当然,如果我们有更大的样本量,我们会包含更多的水果子类别(green_fruits,blue_fruits等).

Answer 1

技巧是建立一个平衡的阵列。我提供了一种笨拙的方法。然后通过引用平衡数组来循环进行 groupby 对象采样。

def rep_sample(df, col, n, *args, **kwargs):
    nu = df[col].nunique()
    m = len(df)
    mpb = n // nu
    mku = n - mpb * nu
    fills = np.zeros(nu)
    fills[:mku] = 1

    sample_sizes = (np.ones(nu) * mpb + fills).astype(int)

    gb = df.groupby(col)

    sample = lambda sub_df, i: sub_df.sample(sample_sizes[i], *args, **kwargs)

    subs = [sample(sub_df, i) for i, (_, sub_df) in enumerate(gb)]

    return pd.concat(subs)

示范

rep_sample(df, 'category1', 5)