如何高效更新Pandas DataFrame各分组内的PRIMARY_INDICATOR值

场景与需求 #

现有包含ID、DATE、PRIMARY_INDICATOR、PHONE列的Pandas DataFrame，每个ID对应多行数据，且每个ID分组内已按DATE降序排列。需完成以下处理：

• 若ID分组内PRIMARY_INDICATOR为1的数量≤1，无需改动；
• 若数量>1，仅保留DATE最新的那行值为1，其余设为0。

当前实现问题 #

原代码通过循环每个ID创建子DataFrame处理，在28万唯一ID的场景下耗时80分钟，核心瓶颈是频繁创建子DataFrame的操作。原代码如下：

import pandas as pd

data = {'ID': [123, 123, 999, 999, 999, 765, 765], 
        'DATE': ['20230125', '20230124', '20230125', '20230120', '20230125', '20230125', '20230125'],
        'PRIMARY_INDICATOR': [1, 0, 1, 1, 0, 0, 0],
        'PHONE' : [8071234, 8079999, 8074312, 9087654, 1235678, 9990000, 9999999]}
df = pd.DataFrame.from_dict(data)
idSet = set(df.ID.unique())
while idSet :
  currentId = idSet.pop()
  idDataframe = df.loc[df['ID'] == currentId]
  idDataframe.drop_duplicates()
  primaryPhoneIndicated = False
  for index, row in idDataframe.iterrows():
    if not primaryPhoneIndicated and row['PRIMARY_INDICATOR'] == 1:
      PRIMARY_INDICATOR = 1
      primaryPhoneIndicated = True
    else:
      PRIMARY_INDICATOR = 0
    print([row['ID'], row['DATE'], PRIMARY_INDICATOR, row['PHONE']])

高效原生Pandas实现方案 #

利用Pandas的groupby结合向量化操作，避免循环创建子DataFrame，大幅提升处理速度。以下是两种可行方案：

方案一：使用groupby+mask+cumsum #

import pandas as pd

data = {'ID': [123, 123, 999, 999, 999, 765, 765], 
        'DATE': ['20230125', '20230124', '20230125', '20230120', '20230125', '20230125', '20230125'],
        'PRIMARY_INDICATOR': [1, 0, 1, 1, 0, 0, 0],
        'PHONE' : [8071234, 8079999, 8074312, 9087654, 1235678, 9990000, 9999999]}
df = pd.DataFrame.from_dict(data)

# 可选：将DATE转为日期类型（若已按降序排列可跳过）
df['DATE'] = pd.to_datetime(df['DATE'], format='%Y%m%d')

# 核心逻辑：分组后，仅保留每个组内第一个PRIMARY_INDICATOR为1的行，其余1转为0
df['PRIMARY_INDICATOR'] = df.groupby('ID').apply(
    lambda x: x['PRIMARY_INDICATOR'].mask(
        (x['PRIMARY_INDICATOR'] == 1) & (x['PRIMARY_INDICATOR'].cumsum() > 1), 0
    )
).reset_index(level=0, drop=True)

print(df)

方案二：使用groupby+自定义处理函数 #

import pandas as pd

data = {'ID': [123, 123, 999, 999, 999, 765, 765], 
        'DATE': ['20230125', '20230124', '20230125', '20230120', '20230125', '20230125', '20230125'],
        'PRIMARY_INDICATOR': [1, 0, 1, 1, 0, 0, 0],
        'PHONE' : [8071234, 8079999, 8074312, 9087654, 1235678, 9990000, 9999999]}
df = pd.DataFrame.from_dict(data)

def process_primary_indicator(group):
    # 找到组内第一个PRIMARY_INDICATOR为1的索引
    first_one_mask = group['PRIMARY_INDICATOR'] == 1
    if first_one_mask.any():
        first_one_idx = group[first_one_mask].index[0]
        # 将除第一个1外的其他1设为0
        group.loc[group.index != first_one_idx, 'PRIMARY_INDICATOR'] = group.loc[group.index != first_one_idx, 'PRIMARY_INDICATOR'].replace(1, 0)
    return group

# 分组应用处理函数
df = df.groupby('ID', group_keys=False).apply(process_primary_indicator)

print(df)

效果说明 #

上述方案均采用Pandas原生的向量化分组操作，避免了循环创建子DataFrame的开销，处理28万唯一ID的场景下，耗时可压缩至数分钟甚至更短，性能提升显著。

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