Python选择N行数数据框

Question

我有一个包含2列的数据框，我想从A列的B列中选择N个行

   A    B
   0    A
   0    B
   0    I
   0    D
   1    A
   1    F
   1    K
   1    L
   2    R

对于A列中的每个唯一数字，请给我B列中的N个随机行：如果N == 2，则结果数据帧将如下所示。如果A列最多不包含N行，则返回A列的全部

   A    B
   0    A
   0    D
   1    F
   1    K
   2    R

Answer 1

对DataFrame.sample中的每个组使用GroupBy.apply，并用if-else测试组的长度：

N = 2
df1 = df.groupby('A').apply(lambda x: x.sample(N) if len(x) >=N else x).reset_index(drop=True)
print (df1)
   A  B
0  0  I
1  0  D
2  1  A
3  1  K
4  2  R

或者：

N = 2
df1 = df.groupby('A', group_keys=False).apply(lambda x: x.sample(N)  if len(x) >=N else x)
print (df1)
   A  B
0  0  A
3  0  D
5  1  F
6  1  K
8  2  R