如何用Pandas内置函数按指定规则为DataFrame分组（无需逐行循环）

问题背景 #

给定如下DataFrame：

import pandas as pd
max_skip=2
data = {'col1':['NT', 'NT', 'NT', 'T','T','T','NT', 'NT', 'T','T','T','NT', 'NT', 'NT','T','T','T','T','NT','NT']}
 
# 创建DataFrame
df = pd.DataFrame(data)

需要按照以下规则对数据分组：

• 若连续NT的行数超过max_skip（值为2），则将这些行分配至组0；否则分配至当前组ID（组ID从0开始）
• 若前一行所在的连续NT行数超过2（max_skip），则下一行的组ID递增；否则保持当前组ID不变

期望得到的结果：

df['group']=[0,0,0,1,1,1,1,1,1,1,1,0,0,0,2,2,2,2,2,2]

解决方案 #

可以通过Pandas的窗口函数、分组标记和累计计算实现，无需逐行循环，步骤如下：

步骤1：标记连续块并计算块长度 #

先识别连续的NT/T块，再计算每个块的行数：

# 给每个连续相同值的块分配唯一ID
df['block_id'] = (df['col1'] != df['col1'].shift()).cumsum()
# 计算每个块的行数，同步到每一行
df['block_len'] = df.groupby('block_id')['col1'].transform('size')

步骤2：标记超长NT块与组ID递增触发点 #

标记需要归为组0的超长NT块，以及需要触发组ID递增的位置：

# 标记当前行是否属于超长NT块（长度>max_skip）
df['is_long_nt'] = (df['col1'] == 'NT') & (df['block_len'] > max_skip)
# 标记下一个块是否需要递增组ID（当前块是超长NT块时，下一个块的第一行触发递增）
df['inc_trigger'] = df['is_long_nt'].shift(-1).fillna(False)

步骤3：计算最终组ID #

先通过触发点累计生成初始组ID，再将超长NT块的组ID重置为0：

# 从0开始，遇到触发点就递增组ID
df['group'] = df['inc_trigger'].cumsum()
# 将超长NT块的组ID强制设为0
df.loc[df['is_long_nt'], 'group'] = 0

完整代码 #

整合所有步骤，最终代码如下：

import pandas as pd
max_skip=2
data = {'col1':['NT', 'NT', 'NT', 'T','T','T','NT', 'NT', 'T','T','T','NT', 'NT', 'NT','T','T','T','T','NT','NT']}
 
# 创建DataFrame
df = pd.DataFrame(data)

# 标记连续块并计算长度
df['block_id'] = (df['col1'] != df['col1'].shift()).cumsum()
df['block_len'] = df.groupby('block_id')['col1'].transform('size')

# 标记超长NT块和递增触发点
df['is_long_nt'] = (df['col1'] == 'NT') & (df['block_len'] > max_skip)
df['inc_trigger'] = df['is_long_nt'].shift(-1).fillna(False)

# 生成组ID并修正超长NT块
df['group'] = df['inc_trigger'].cumsum()
df.loc[df['is_long_nt'], 'group'] = 0

# 可选：删除中间辅助列
df = df.drop(['block_id', 'block_len', 'is_long_nt', 'inc_trigger'], axis=1)

print(df)

运行后即可得到符合要求的分组结果。

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