pandas 中的随机分层抽样

Question

我创建了一个 pandas 数据框，如下所示：

import pandas as pd
import numpy as np

ds = {'col1' : [1,1,1,1,1,1,1,2,2,2,2,3,3,3,3,3,4,4,4,4,4,4,4,4,4],
      'col2' : [12,3,4,5,4,3,2,3,4,6,7,8,3,3,65,4,3,2,32,1,2,3,4,5,32],
      }

df = pd.DataFrame(data=ds)

数据框如下所示：

print(df)

    col1  col2
0      1    12
1      1     3
2      1     4
3      1     5
4      1     4
5      1     3
6      1     2
7      2     3
8      2     4
9      2     6
10     2     7
11     3     8
12     3     3
13     3     3
14     3    65
15     3     4
16     4     3
17     4     2
18     4    32
19     4     1
20     4     2
21     4     3
22     4     4
23     4     5
24     4    32

根据

col1

列的值，我需要提取：

• 3 条随机记录，其中

  col1 == 1

• 2 条随机记录，使得

  col1 = 2

• 2 条随机记录，使得

  col1 = 3

• 3 条随机记录，使得

  col1 = 4

有人可以帮助我吗？

Answer 1

我会用

sample(frac=1)

打乱输入方式，然后计算

groupby.cumcount

来选择每组的第一个

N

样本，其中

N

在字典中定义：

# {col1: number of samples}
n = {1: 3, 2: 2, 3: 2, 4: 3}

out = df[df[['col1']].sample(frac=1)
                     .groupby('col1').cumcount()
                     .lt(df['col1'].map(n))]

输出示例：

    col1  col2
0      1    12
3      1     5
4      1     4
7      2     3
8      2     4
11     3     8
13     3     3
17     4     2
18     4    32
24     4    32