在Pandas中按条件为同policyId添加统一的contract_type列

问题描述 #

现有如下DataFrame：

import pandas as pd

data = {
    'policyId': ['X','X','X','X', 'Y','Y','Y','Y', 'Z', 'Z','Z','Z'],
    'element_id': [242, 243, 241, 257, 242, 243, 241, 257, 242, 243, 241, 257],
    'element_selected': [True, True, True, True, False, True, True, True, True, True, True, False]
}
df = pd.DataFrame(data)

所有policyId包含相同的element_id，差异仅在于element_id是否被选中（由element_selected字段标识）。需要新增contract_type列，规则如下：

• 若element_id=242被选中，contract_type为"CONTENT"
• 若element_id=257被选中，contract_type为"PRIVATE LIABILITY"
• 若两者均被选中，则为"CONTENT + PRIVATE LIABILITY"
  要求每个policyId对应的contract_type为统一值。

预期输出：

policyId    element_id  element_selected    contract_type
X                  242  True                CONTENT + PRIVATE LIABILITY
X                  243  True                CONTENT + PRIVATE LIABILITY
X                  241  True                CONTENT + PRIVATE LIABILITY
X                  257  True                CONTENT + PRIVATE LIABILITY
Y                  242  False               PRIVATE LIABILITY
Y                  243  True                PRIVATE LIABILITY
Y                  241  True                PRIVATE LIABILITY
Y                  257  True                PRIVATE LIABILITY
Z                  242  True                CONTENT
Z                  243  True                CONTENT
Z                  241  True                CONTENT
Z                  257  False               CONTENT

解决方案 #

通过分组聚合+合并的方式实现，步骤如下：

1. 提取每个policyId对应242和257的选中状态：

# 透视数据，得到每个policyId下242、257的选中状态
pivot_df = df.pivot(index='policyId', columns='element_id', values='element_selected').reset_index()
contract_info = pivot_df[['policyId', 242, 257]].rename(columns={242: 'has_content', 257: 'has_private'})

2. 根据状态生成统一的contract_type：

# 按规则拼接合同类型
def generate_contract(row):
    type_list = []
    if row['has_content']:
        type_list.append("CONTENT")
    if row['has_private']:
        type_list.append("PRIVATE LIABILITY")
    return " + ".join(type_list)

contract_info['contract_type'] = contract_info.apply(generate_contract, axis=1)

3. 将结果合并回原DataFrame：

# 合并后每个policyId的所有行都会带上统一的contract_type
result_df = df.merge(contract_info[['policyId', 'contract_type']], on='policyId', how='left')

执行上述代码后，result_df即为符合预期的输出。

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