Pandas：为DataFrame每行生成右侧全为1的基准年份列

现有数据集 #

import pandas as pd
import numpy as np

data = {'ID': ['A', 'B', 'C', 'D'], 
        '2012': [0, 1, 1, 1], 
        '2013': [0, 0, 1, 1], 
        '2014': [0, 0, 0, 1], 
        '2015': [0, 0, 1, 1], 
        '2016': [0, 0, 1, 0], 
        '2017': [1, 0, 1, 1]}

df = pd.DataFrame(data)

需求说明 #

为每行新增列Baseline_Year，取值规则：

• 找到该行中最左侧的年份列，要求该列值为1，且其右侧所有年份列的值全为1
• 若不存在符合条件的列，设为缺失值（np.nan）

预期结果 #

最终输出的DataFrame如下（新增ID为E的行以覆盖全1场景）：

data = {'ID': ['A', 'B', 'C', 'D', 'E'], 
        '2012': [0, 1, 1, 1, 1], 
        '2013': [0, 0, 1, 1, 1], 
        '2014': [0, 0, 0, 1, 1], 
        '2015': [0, 0, 1, 1, 1], 
        '2016': [0, 0, 1, 0, 1], 
        '2017': [1, 0, 1, 1, 1],
        'Baseline_Year': [np.nan, np.nan, '2015','2017', '2012'],
       }

df_results = pd.DataFrame(data)
print(df_results)

输出内容：

ID  2012  2013  2014  2015  2016  2017 Baseline_Year
0  A     0     0     0     0     0     1            NaN
1  B     1     0     0     0     0     0            NaN
2  C     1     1     0     1     1     1          2015
3  D     1     1     1     1     0     1          2017
4  E     1     1     1     1     1     1          2012

解决方案代码 #

# 提取所有年份列（排除ID列）
year_cols = df.columns[1:]

def find_baseline_year(row):
    # 从左到右遍历年份列，找第一个符合条件的列
    for col in year_cols:
        col_idx = year_cols.get_loc(col)
        # 获取当前列右侧的所有列
        right_cols = year_cols[col_idx+1:]
        
        # 验证条件：当前列值为1，且右侧所有列值全为1（无右侧列时默认满足）
        if row[col] == 1 and (right_cols.empty or (row[right_cols] == 1).all()):
            return col
    # 无符合条件的列返回缺失值
    return np.nan

# 生成新列
df['Baseline_Year'] = df.apply(find_baseline_year, axis=1)

# 添加测试行E
df = pd.concat([df, pd.DataFrame({'ID': ['E'], '2012': [1], '2013': [1], '2014': [1], '2015': [1], '2016': [1], '2017': [1]})], ignore_index=True)

print(df)

代码说明 #

1. 提取年份列：分离出DataFrame中用于判断的年份列集合，排除ID列。
2. 遍历判断逻辑：对每行从左到右检查年份列，验证当前列值为1且右侧所有列值全为1（无右侧列时默认符合条件），找到第一个符合条件的列名返回。
3. 生成新列：通过apply方法将判断逻辑应用到每行，生成Baseline_Year列。

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