Python Pandas数据处理：按P_ID生成New_column逻辑问题求助

现有如下DataFrame：

import pandas as pd
df = pd.DataFrame({
    'P_ID': [348374, 348374, 348374, 1348374, 1348374, 1348374, 1348374],
    'I_ISSUE_TYPE_CODE': ['CLAIM', 'RSA', 'RSA', 'EPQR', 'CLAIM', 'EPQR', 'EPQR'],
    'A_b': [234, 5671, 8900, 3451, 660, 770, 660]
})

需要实现以下逻辑：

• 按P_ID分组，若组内存在I_ISSUE_TYPE_CODE为RSA的记录，仅首次出现的RSA对应的New_column设为1，其余行设为0；
• 若组内无RSA记录，则组内首行的New_column设为1，其余行设为0。

预期输出如下：

df = pd.DataFrame({
    'P_ID': [348374, 348374, 348374, 1348374, 1348374, 1348374, 1348374],
    'I_ISSUE_TYPE_CODE': ['CLAIM', 'RSA', 'RSA', 'EPQR', 'CLAIM', 'EPQR', 'EPQR'],
    'A_b': [234, 5671, 8900, 3451, 660, 770, 660],
    'New_column': [0, 1, 0, 1, 0, 0, 0]
})

尝试了以下代码，但未得到正确结果：

df['New_column'] = df.groupby('P_ID')['I_ISSUE_TYPE_CODE'].apply(lambda x: (x == 'RSA').cumsum().eq(1).astype(int))

之前的代码仅处理了RSA首次出现的场景，但未覆盖无RSA组需首行设1的需求，以下是两种可行的实现方式：

方法一：自定义分组处理函数 #

通过groupby.apply对每个分组单独判断处理：

def process_group(group):
    # 检查当前分组是否包含RSA
    has_rsa = (group['I_ISSUE_TYPE_CODE'] == 'RSA').any()
    group['New_column'] = 0
    if has_rsa:
        # 找到首次出现RSA的索引并设为1
        first_rsa_pos = (group['I_ISSUE_TYPE_CODE'] == 'RSA').idxmax()
        group.loc[first_rsa_pos, 'New_column'] = 1
    else:
        # 无RSA则将分组首行设为1
        group.iloc[0, group.columns.get_loc('New_column')] = 1
    return group

# 应用分组处理，group_keys=False避免保留分组键索引
df = df.groupby('P_ID', group_keys=False).apply(process_group)

方法二：基于transform的简洁写法 #

通过两个辅助标记列合并条件，最后生成目标列：

# 标记首次出现RSA的行
df['first_rsa'] = df.groupby('P_ID')['I_ISSUE_TYPE_CODE'].transform(
    lambda x: (x == 'RSA') & (x.cumsum() == 1)
)
# 标记无RSA分组的首行
df['first_row_no_rsa'] = df.groupby('P_ID')['I_ISSUE_TYPE_CODE'].transform(
    lambda x: (x != 'RSA').all() & (x.reset_index(drop=True).index == 0)
)
# 合并两个条件得到New_column
df['New_column'] = (df['first_rsa'] | df['first_row_no_rsa']).astype(int)
# 清理辅助列（可选）
df = df.drop(['first_rsa', 'first_row_no_rsa'], axis=1)

两种方法都能得到符合预期的结果，可根据代码风格偏好选择。

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