如何在pandas DataFrame中选择不同条件下的前N个主题？

Question

我想从各个儿童组下的每个家长组中选择得分最高的 N & M 科目。在这里，我想从parent_group = A，child_group = 1，parent_group = B，child_group = 1和parent_group = C，child_group = 1中选择前2个科目。重复相同的操作，为child_group = 2选择前3个科目。

import pandas as pd 
import numpy as np 

n = 300
np.random.seed(1)
parent_groups = np.random.choice(['A','B','C'], n)
child_groups = np.random.choice([1,2], n)
scores = np.random.random(size=n)
subjects = np.random.choice(range(100, 2000), n)

df = pd.DataFrame({'subjects':subjects,
                   'parent_groups':parent_groups,
                   'child_groups':child_groups,
                   'scores':scores})
df

    subjects    parent_groups   child_groups    scores
0   559         B                 1             0.969565
1   950         A                 2             0.303878
2   400         A                 2             0.313755
3   1956        B                 2             0.226462
4   1673        B                 2             0.672335
... ... ... ... ...

我计划根据子组将 df 分为两部分，例如，

df_1 = df[df['child_groups'] == 1]
df_2 = df[df['child_groups'] == 2]
df_1


    subjects    parent_groups   child_groups    scores
0   559         B               1               0.969565
5   1299        A               1               0.778922
7   574         B               1               0.615265
9   1907        B               1               0.446100
10  549         A               1               0.968013
... ... ... ... ...

然后根据最高分为每个家长组选择前 2 个科目。对 df_2 重复相同的过程，选择前 3 个主题。最后将这些受试者及其分数、parent_groups 和 child_groups 信息组合成最终的数据帧。但是，我不知道如何以任何有效的方式进行此操作。有什么建议么？谢谢

Answer 1

您可以按

child_groups

和

parent_groups

，然后按

apply

nlargest

进行分组，使用

child_groups

数字选择 2 或 3 个最大

scores

值：

out = (df
    .groupby(['child_groups', 'parent_groups'], as_index=False)
    .apply(lambda g:g.nlargest(1+g['child_groups'].iloc[0], 'scores'))
    .reset_index(drop=True)
)

输出：

    subjects parent_groups  child_groups    scores
0       1338             A             1  0.986943
1        566             A             1  0.976481
2       1911             B             1  0.992827
3       1843             B             1  0.980412
4        820             C             1  0.993938
5        141             C             1  0.932257
6       1624             A             2  0.994365
7        962             A             2  0.928267
8        846             A             2  0.886920
9        963             B             2  0.974991
10      1819             B             2  0.921134
11      1446             B             2  0.879272
12      1900             C             2  0.999769
13       480             C             2  0.991654
14      1999             C             2  0.991451