按UID筛选首个满足后续年份全为Good?=1的年份并生成新列

问题背景 #

原始数据集 #

import pandas as pd

df = pd.DataFrame({
    'UID': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'Year': [2015, 2016, 2017, 2014, 2015, 2017, 2014, 2015, 2016],
    'Good?': [0, 1, 1, 0, 0, 1, 0, 1, 0]
})

需求说明 #

对每个UID，找到第一个Good?值为1且后续所有年份的Good?值均为1的Year值；若不满足条件（包括没有任何Good?=1的情况），则赋值为2017。

出错的现有代码 #

# Group the DataFrame by UID
groups = df.groupby('UID')

# Initialize an empty list to store the results
results = []

# Loop over each UID group
for uid, group in groups:
    # Find the first index with a Good value of 1
    first_good_index = group[group['Good?'] == 1].index[0]
    print(first_good_index)
    
    # Check if all following years have a Good value of 1
    if (group.loc[first_good_index+1:, 'Good?'] == 1).all():
        # If so, append the UID and the year of the first good row to the results list
        results.append((uid, group.loc[first_good_index, 'Year']))
    else:
        results.append((uid, 2017))

# Create a DataFrame from the results
results_df = pd.DataFrame(results, columns=['UID', 'First Good Year'])

# Print the results
print(results_df)

预期结果 #

results_df = pd.DataFrame({
    'UID': [1, 2, 3],
    'First Good Year': [2016, 2017, 2017],
})

print(results_df)
# 输出：
#    UID  First Good Year
# 0    1             2016
# 1    2             2017
# 2    3             2017

错误原因分析 #

1. 无匹配项时的索引错误：当某个UID没有任何Good?=1的记录时，group[group['Good?'] ==1].index[0]会抛出IndexError，因为空的Series无法通过[0]取值。
2. 依赖全局索引的风险：原代码使用DataFrame的全局索引进行切片，若组内索引不连续或不是从0开始，会导致切片逻辑错误；另外当第一个Good?=1的行是组内最后一行时，first_good_index+1:会生成空切片，虽然.all()返回True，但逻辑依赖全局索引不够稳健。

修正后的代码 #

import pandas as pd

df = pd.DataFrame({
    'UID': [1, 1, 1, 2, 2, 2, 3, 3, 3],
    'Year': [2015, 2016, 2017, 2014, 2015, 2017, 2014, 2015, 2016],
    'Good?': [0, 1, 1, 0, 0, 1, 0, 1, 0]
})

# 先按UID和Year排序，确保每个组内年份升序
df_sorted = df.sort_values(['UID', 'Year']).reset_index(drop=True)

results = []
for uid, group in df_sorted.groupby('UID'):
    # 获取组内Good?=1的行的位置（相对组内的索引）
    good_positions = group[group['Good?'] == 1].index
    if not good_positions.empty:
        first_good_pos = good_positions[0]
        # 检查第一个Good之后的所有行是否全为1
        # 使用组内的iloc切片，避免全局索引问题
        if group.iloc[first_good_pos+1:]['Good?'].eq(1).all():
            results.append((uid, group.loc[first_good_pos, 'Year']))
        else:
            results.append((uid, 2017))
    else:
        # 没有Good?=1的记录，直接赋值2017
        results.append((uid, 2017))

results_df = pd.DataFrame(results, columns=['UID', 'First Good Year'])
print(results_df)

代码说明 #

• 排序与重置索引：先对数据集按UID和Year排序，确保每个组内年份是升序排列，同时重置索引避免全局索引干扰。
• 处理无匹配项的情况：通过good_positions.empty判断是否存在Good?=1的记录，避免索引错误。
• 组内相对位置切片：使用iloc进行组内切片，基于组内的相对位置而非全局索引，逻辑更稳健。
• 空切片的兼容：当第一个Good?=1的行是组内最后一行时，group.iloc[first_good_pos+1:]会返回空DataFrame，.all()会返回True，符合需求逻辑。

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