Pandas标记非唯一行分组 适配排除全Structure=s的配对规则

场景说明 #

这是StackOverflow相关问题的扩展场景，需要实现纯pandas向量化（无显式循环）的行标记逻辑。

规则要求 #

• 分组边界：以TDateID、GroupID的组合作为一级分组范围，在范围内识别非唯一行分组
• ID规则：为所有分组分配全局连续自增整型ID，同组行共享相同ID
• 特殊约束：如果某一级分组内所有行的Structure字段值均为"s"，则不执行跨子组配对，组内每行分配独立自增ID
• 配对逻辑：仅当一级分组内存在非"s"值时，按行位序对不同Cents、SD_YF的子组做配对，同位序的行归入同一组共享ID

测试数据集构造 #

import pandas as pd
df = pd.DataFrame({'Index': [10, 11, 12, 13, 14, 15, 4, 5, 21, 22, 24, 25],
 'Cents': [182.5,
  182.5,
  153.5,
  153.5,
  43.0,
  43.0,
  152.0,
  152.0,
  53.0,
  53.0,
  252.0,
  252.0],
 'Structure': ['s', 's', 's', 's', 'p', 'p', 's', 's', 's', 's', 's', 's'],
 'SD_YF': [2.1,
  2.1,
  1.05,
  1.05,
  11.0,
  11.0,
  21.0,
  21.0,
  13.0,
  13.0,
  25.0,
  25.0],
 'TDateID': [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2],
 'GroupID': [0, 0, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]})

预期效果 #

新增UniID自增字段：

• 例如TDateID=1、GroupID=2下，Structure=p的两行与Structure=s的两行按行位序配对，分别共享UniID=5、UniID=6
• 其余所有Structure全为s的重复行，每行分配独立的自增ID

原有实现的缺陷 #

以下基础实现可完成常规分组标记，但无法适配「分组内存在非s值才执行配对」的规则：

df['UniID'] = (df['GroupID']
 +df.groupby('GroupID').ngroup().add(1)
 +df.groupby(['GroupID', 'Cents', 'SD_YF']).cumcount()
)

适配规则的向量化实现 #

完全无显式循环，全pandas内置向量化操作实现：

# 标记一级分组是否需要执行配对逻辑（组内存在非s值则需要配对）
df['_pair_flag'] = df.groupby(['TDateID', 'GroupID'])['Structure'].transform(lambda x: x.ne('s').any())

# 分别计算两种场景下的组内子ID
# 需配对场景：按(Cents, SD_YF)分子组取累计计数，同计数即为同位序配对组
pair_sub = df.groupby(['TDateID', 'GroupID', 'Cents', 'SD_YF']).cumcount()
# 不需配对场景（全s组）：组内每行独立，直接取组内行序号作为子ID
nopair_sub = df.groupby(['TDateID', 'GroupID']).cumcount()

# 按标记拼接子ID，再转换为全局连续自增ID
df['sub_id'] = pair_sub.where(df['_pair_flag'], nopair_sub)
df['UniID'] = df.groupby(['TDateID', 'GroupID', 'sub_id']).ngroup() + 1

# 删除中间辅助列
df.drop(columns=['_pair_flag', 'sub_id'], inplace=True)

执行后输出结果完全匹配预期：全s的一级分组内每行ID唯一，存在非s值的一级分组内按位序配对共享ID，ID全局连续自增。

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