如何用Python按分组ID及Cycle阈值设置Salary列的0/1值？

问题描述 #

现有如下Python数据集并已转换为DataFrame：

import pandas as pd

# Create dataset
data = {'id': [1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
        'cycle': [1, 2, 3, 4, 5,6,7,8,9, 1, 2, 3,4,5,6, 1, 2, 3, 4,5,6,7,8],
        'Salary': [0, 0, 0,  0,0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        'Days': [123, 128, 66, 120, 141, 128, 66, 120, 141, 52,96, 120, 141, 52,96, 120, 141,123,15,85,36,58,89],
        }

# Convert to dataframe
df = pd.DataFrame(data)

需要针对每个id分组，根据Cycle阈值设置Salary列值：

• id=1时，cycle>=4则Salary设为1
• id=2时，cycle>=3则Salary设为1
• id=3时，cycle>=6则Salary设为1

实现方法 #

方法1：字典映射 + 布尔索引 #

这种方法高效直观，适合简单的映射逻辑：

# 定义id与cycle阈值的映射关系
threshold_map = {1:4, 2:3, 3:6}

# 遍历每个id及其阈值，批量赋值
for id_val, threshold in threshold_map.items():
    df.loc[(df['id'] == id_val) & (df['cycle'] >= threshold), 'Salary'] = 1

# 查看结果
print(df)

方法2：使用np.select多条件判断 #

可读性强，适合条件较多的场景：

import numpy as np

# 定义条件列表和对应赋值
conditions = [
    (df['id'] == 1) & (df['cycle'] >= 4),
    (df['id'] == 2) & (df['cycle'] >= 3),
    (df['id'] == 3) & (df['cycle'] >= 6)
]
values = [1, 1, 1]

# 应用条件，未满足条件的行保留原Salary值
df['Salary'] = np.select(conditions, values, default=df['Salary'])

方法3：分组后自定义函数处理 #

扩展性强，适合后续有更复杂分组逻辑的场景：

def set_salary(group):
    id_val = group['id'].iloc[0]
    if id_val == 1:
        group.loc[group['cycle'] >=4, 'Salary'] =1
    elif id_val ==2:
        group.loc[group['cycle'] >=3, 'Salary'] =1
    elif id_val ==3:
        group.loc[group['cycle'] >=6, 'Salary'] =1
    return group

df = df.groupby('id').apply(set_salary).reset_index(drop=True)

以上三种方法均可实现需求，其中方法1在数据量较大时性能最优，方法2和3更适配复杂逻辑场景。

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