Pandas Python选择N行数数据帧

我有一个有2列的数据框，我想从B列中为a列选择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

在中使用每组，测试组的长度为

如果其他

：

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