基于fips-year唯一组合添加nonattainment_pm列的实现方案

我有如下结构的DataFrame：

补充信息 #

• nonattainment列的取值为P或W
• pollutant列中的PM类污染物为：['PM-2.5 (1997)', 'PM-2.5 (2006)', 'PM-10 (1987)', 'PM-2.5 (2012)']

任务需求 #

需要添加名为nonattainment_pm的新列：对于每个唯一的fips-year组合，若该组合内存在属于上述PM类的污染物，且对应nonattainment值为P或W，则该组合下所有行的nonattainment_pm值为1，否则为空。

预期输出 #

新的DataFrame结构如下：

可以通过分组计算+映射的方式实现，步骤如下：

1. 定义PM类污染物列表
2. 标记每行是否符合条件（是PM类污染物且nonattainment为P/W）
3. 按fips和year分组，判断每组是否存在符合条件的行
4. 将分组结果映射回原DataFrame，生成新列

代码实现：

import pandas as pd

# 定义PM类污染物列表
pm_pollutants = ['PM-2.5 (1997)', 'PM-2.5 (2006)', 'PM-10 (1987)', 'PM-2.5 (2012)']

# 示例数据（可替换为你的实际DataFrame）
data = {
    'fips': [72137]*8,
    'year': [1992]*4 + [1993]*4,
    'pollutant': [
        'Sulfur Dioxide (1971)',
        'PM-2.5 (1997)',
        '8-Hour Ozone (2015)',
        "'Nitrogen Dioxide (1971)'",
        'Sulfur Dioxide (1971)',
        'PM-2.5 (1997)',
        '8-Hour Ozone (2015)',
        "'Nitrogen Dioxide (1971)'"
    ],
    'nonattainment': ['', 'P', 'W', '', '', '', 'W', '']
}
df = pd.DataFrame(data)

# 标记符合条件的行
df['is_pm_eligible'] = df.apply(
    lambda row: row['pollutant'] in pm_pollutants and row['nonattainment'] in ['P', 'W'],
    axis=1
)

# 按fips-year分组，判断每组是否存在符合条件的行
group_result = df.groupby(['fips', 'year'])['is_pm_eligible'].any().reset_index()
group_result['nonattainment_pm'] = group_result['is_pm_eligible'].map({True: 1, False: None})

# 将结果合并回原DataFrame
df = df.merge(group_result[['fips', 'year', 'nonattainment_pm']], on=['fips', 'year'], how='left')

# 移除临时列（可选）
df = df.drop('is_pm_eligible', axis=1)

print(df)

代码说明 #

• is_pm_eligible是临时标记列，用于判断当前行是否满足“PM类污染物+nonattainment为P/W”的条件
• 分组后用any()判断每组是否存在至少一行符合条件，只要该组合内有一行满足，整个组的nonattainment_pm都设为1
• 通过merge将分组结果映射回原DataFrame，保证每个fips-year组合下的所有行有相同的nonattainment_pm值

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