Python按组筛选列中首次出现符号变化的行

需求说明 #

给定如下DataFrame：

col1 col2 group
10   1     A
11   2     A
12  -2     A
13  -3     B
14   3     B

需要按group分组，提取每个组内col2首次出现符号变化的行，最终结果如下：

col1 col2 group
12  -2    A
14   3    B

实现步骤与代码 #

1. 初始化数据

首先导入依赖库并创建DataFrame：

import pandas as pd
import numpy as np

data = {
    'col1': [10, 11, 12, 13, 14], 
    'col2': [1, 2, -2, -3, 3],
    'group': ['A', 'A', 'A', 'B', 'B']
}  
df = pd.DataFrame(data)

2. 核心处理逻辑

通过符号标记、分组检测、筛选首行三个步骤完成需求：

# 1. 计算col2的符号（正数=1，负数=-1，0=0）
df['sign'] = np.sign(df['col2'])

# 2. 分组检测符号变化：当前行符号与前一行不同则标记为True
df['is_change'] = df.groupby('group')['sign'].transform(lambda x: x != x.shift())

# 3. 筛选每个分组中第一个符号变化的行
# 排除分组第一行（无前置行，不算变化），再按组取第一个变化行
result = df[(df['is_change']) & (~df.index.isin(df.groupby('group').head(1).index))]
result = result.groupby('group').first().reset_index()[['col1', 'col2', 'group']]

3. 查看结果

执行print(result)后输出：

col1  col2 group
0    12    -2     A
1    14     3     B

简化写法 #

如果追求代码简洁，可以用groupby.apply直接处理每个分组：

df['sign'] = np.sign(df['col2'])

result = df.groupby('group').apply(
    lambda g: g[g['sign'] != g['sign'].shift()].iloc[0] if len(g[g['sign'] != g['sign'].shift()]) > 1 else None
).dropna().reset_index(drop=True)[['col1', 'col2', 'group']]

这里判断len(...)>1是因为分组第一行的shift()会产生NaN，对应的sign != shift()会是True，需要排除这个无效的"变化"。

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