如何根据一列的值有条件地选择前N个组并按两列分组？

Question

这是此帖子

的后续内容

这是我的数据框：

df = pd.DataFrame(
    {
        'a': [10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60, 65, 70, 10, 22],
        'b': [1, 1, 1, -1, -1, -1, -1, 2, 2, 2, 2, -1, -1, -1, -1],
        'c': [25, 25, 25, 45, 45, 45, 45, 65, 65, 65, 65, 40, 40, 30, 30],
        'main': ['x', 'x', 'x', 'x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'y', 'y', 'y', 'y']
    }
)

预期输出：Groupby

main

AND

c

:

    a  b   c main
0   10  1  25    x
1   15  1  25    x
2   20  1  25    x
3   25 -1  45    x
4   30 -1  45    x
5   35 -1  45    x
6   40 -1  45    x
11  65 -1  40    y
12  70 -1  40    y
13  10 -1  30    y
14  22 -1  30    y

过程如下：注意

groupby

是由两列完成的：

所以对于每个

main

：

a) 选择所有

b

值为

1

的组。在我的数据和这个

df

中，只有一组有这种情况。

b) 选择前两组（从

df

顶部开始），它们的所有

b

值均为 -1。

请注意，我的数据中可能不存在具有

a

或

b

条件的组。如果是这种情况，返回符合条件的任何内容都可以。例如，输出可能只有一组或根本没有组。

我想要的群组如下所示：

这是我基于此答案的尝试，但似乎其他一些事情必须改变：

# identify groups with all 1
m1 = df['b'].eq(1).groupby(df['c', 'main']).transform('all')
# identify groups with all -1
m2 = df['b'].eq(-1).groupby(df['c', 'main']).transform('all')
# keep rows of first 2 groups with all -1
m3 = df[['c', 'main']].isin(df.loc[m2, ['c', 'main']].unique()[:2])

# select m1 OR m3
out = df[m1 | m3]

Answer 1

您可以更新之前的代码以获得每个主函数的前 2 个唯一的“c”：

groups = [df['c'], df['main']]
# identify groups with all 1
m1 = df['b'].eq(1).groupby(groups).transform('all')
# identify groups with all -1
m2 = df['b'].eq(-1).groupby(groups).transform('all')
# keep rows of first 2 groups with all -1, per main
keep = set.union(*df.loc[m2, ['c', 'main']].groupby('main')['c']
                    .agg(lambda x: set(x.unique()[:2])))
# {25}
m3 = df['c'].isin(keep)

# select m1 OR m3
out = df[m1 | m3]

输出：

     a  b   c main
0   10  1  25    x
1   15  1  25    x
2   20  1  25    x
3   25 -1  45    x
4   30 -1  45    x
5   35 -1  45    x
6   40 -1  45    x
11  65 -1  40    y
12  70 -1  40    y
13  10 -1  30    y
14  22 -1  30    y