根据样本量随机抽取组内样本

Question

我有一个非常大的数据集，其结构与此类似：

df = pd.DataFrame({
    'ID': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'Group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
    'SampleSize': [4, 4, 4, 4, 4, 4, 1, 1, 1, 1]
})

这意味着，例如，在

Group

1

中，有 6 个不同的单位可供选择（

ID

s），而对于这个

Group

1

，需要选择 4 个单位来形成样本。所以最终我想得到一个额外的列来指示随机选择的样本，如下所示：

df = pd.DataFrame({
    'ID': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'Group': [1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
    'SampleSize': [4, 4, 4, 4, 4, 4, 1, 1, 1, 1]
    'Sample': [0, 1, 1, 1, 0, 1, 0, 1, 0, 0]
})

我尝试过这样的事情：

def select_random_ids(group):
    sample_size = group['SampleSize'].iloc[0]
    selected_ids = np.random.choice(group['ID'], size=sample_size, replace=False)
    return pd.DataFrame({'ID': selected_ids})

和

.apply(select_random_ids))

但我无法让它工作..

Answer 1

尝试：

def binary_mask_random(r, n):
    a = np.zeros(r, dtype="uint8")
    indices = np.random.choice(range(r), size=n, replace=False)
    a[indices] = 1
    return a


df["Sample"] = df.groupby("Group")["SampleSize"].transform(
    lambda size: binary_mask_random(len(size), size.iat[0])
)
print(df)

打印（例如）：

   ID  Group  SampleSize  Sample
0   1      1           4       1
1   2      1           4       0
2   3      1           4       1
3   4      1           4       1
4   5      1           4       0
5   6      1           4       1
6   7      2           1       0
7   8      2           1       1
8   9      2           1       0
9  10      2           1       0