pandas - 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

标签： pandasdataframe

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