如何按两列groupby并基于列值条件筛选目标分组？

原始数据 #

import pandas as pd

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']
    }
)

预期输出 #

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

处理逻辑 #

需按main和c两列分组，对每个main分组执行以下操作：

• 筛选所有b值均为1的分组
• 筛选从DataFrame顶部开始的前两个所有b值均为-1的分组

注：若不存在符合条件的分组，返回所有匹配到的结果即可（比如仅1个或无分组）

修正后的代码 #

原代码未按main分组单独处理前两个-1分组，且分组索引写法有误，以下是修正后的实现：

import pandas as pd

# 步骤1：标记所有b全为1的分组
m1 = df['b'].eq(1).groupby([df['main'], df['c']]).transform('all')

# 步骤2：筛选出所有b全为-1的分组，并按main分组取前2个符合条件的组
valid_neg_groups = df[df['b'].eq(-1).groupby([df['main'], df['c']]).transform('all')]
# 按main分组，提取每个main下的前2个符合条件的c值
top_neg_groups = valid_neg_groups.groupby('main')['c'].unique().apply(lambda x: x[:2]).explode()
# 转换为便于匹配的(main, c)元组集合
top_neg_tuples = set(zip(top_neg_groups.index, top_neg_groups.values))

# 步骤3：标记属于前2个b全为-1分组的行
m3 = df.apply(lambda row: (row['main'], row['c']) in top_neg_tuples, axis=1)

# 步骤4：合并筛选条件，得到最终结果
out = df[m1 | m3]
print(out)

代码解释 #

1. 标记全1分组：通过groupby([main, c])分组后，用transform('all')判断每组是否所有b值都是1，得到布尔序列m1。
2. 筛选前2个全-1分组：
   • 先筛选出所有b全为-1的行，再按main分组提取每个main下的c值，取前2个；
   • 将这些(main, c)组合转为集合，方便后续行匹配。
3. 标记前2个全-1分组的行：遍历每行，判断其(main, c)是否在目标集合中，得到布尔序列m3。
4. 合并结果：取m1或m3为True的行，即为最终筛选结果。

内容的提问来源于stack exchange，提问作者AmirX