随机抽样组:如何指定n，而不是权重？(使用DataFrameGroupBy.sample)

Question

这个问题遵循这个问题 (其他贡献者要求我将其作为一个新问题发布)。

我们有一个模拟df：

df = pd.DataFrame({
        'id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        'country': ['USA', 'USA', 'USA', 'USA', 'USA', 'Canada', 'Canada', 'Canada', 'USA', 'Canada']
})

假设我想从美国采集4行随机数据，从加拿大采集2行随机数据。我试过：

df.groupby("country").sample(n=[4, 2])

这将返回一个错误。错误可能是使用方括号。那么，如何为每个组指定不同的n呢？

注意，理想情况下，我需要一个使用df.groupby.sample的解决方案。还要注意，我需要指定n，而不是文档中的比例或权重(参见这里)。最后，请注意，我还需要设置一个种子。谢谢

Answer 1

您可以在group上country上获取数据，然后分别从字典中获取要采集的样本数的.sample每组，最后对所有抽样组进行.concat：

d = {'USA': 4, 'Canada': 2} # mapping dict
pd.concat([g.sample(d[k]) for k, g in df.groupby('country', sort=False)])

   id country
0   1     USA
4   5     USA
1   2     USA
2   3     USA
6   7  Canada
9  10  Canada