基于flag列生成DataFrame新列Final的Pandas实现问题

样本数据 #

import pandas as pd

df = pd.DataFrame({
    'A_flag': [1, 1, 1], 
    'B_flag': [1, 1, 0],
    'C_flag': [0, 1, 0],
    'A_value': [5, 3, 7], 
    'B_value': [2, 7, 4],
    'C_value': [4, 2, 5]
})

# 目标结果示例
df1 = pd.DataFrame({
    'A_flag': [1, 1, 1], 
    'B_flag': [1, 1, 0],
    'C_flag': [0, 1, 0],
    'A_value': [5, 3, 7], 
    'B_value': [2, 7, 4],
    'C_value': [4, 2, 5], 
    'Final': [3.5, 3, 7]
})

需求规则 #

需要基于A_flag、B_flag、C_flag三列生成Final新列，规则如下：

• (a) 若三列中等于1的数量为3，Final等于A_value、B_value、C_value的中位数；
• (b) 若等于1的数量为2，Final等于对应两个value列的平均值；
• (c) 若等于1的数量为1，Final等于对应那一个value列的值。

尝试的代码 #

df.loc[df[['A_flag','B_flag','C_flag']].eq(1).sum(axis=1)==3, "Final"]= df[['A_value','B_value','C_value']].median(axis=1)

df.loc[df[['A_flag','B_flag','C_flag']].eq(1).sum(axis=1)==2, "Final"]=
df.loc[df[['A_flag','B_flag','C_flag']].eq(1).sum(axis=1)==1, "Final"]=  

解决方案 #

方法1：向量化处理（高效推荐） #

通过将flag为0的对应value值设为NaN，再根据每行非NaN值的数量计算结果：

import pandas as pd

df = pd.DataFrame({
    'A_flag': [1, 1, 1], 
    'B_flag': [1, 1, 0],
    'C_flag': [0, 1, 0],
    'A_value': [5, 3, 7], 
    'B_value': [2, 7, 4],
    'C_value': [4, 2, 5]
})

flag_cols = ['A_flag', 'B_flag', 'C_flag']
value_cols = ['A_value', 'B_value', 'C_value']

# 筛选出flag为1对应的value，其余设为NaN
masked_values = df[value_cols].where(df[flag_cols].eq(1))

# 按规则计算Final列
df['Final'] = masked_values.apply(
    lambda row: row.median() if row.count() == 3 
              else row.mean() if row.count() == 2 
              else row.dropna().iloc[0],
    axis=1
)

print(df)

方法2：逐行处理（直观易读） #

写一个自定义函数，逐行筛选flag为1的value后计算：

def get_final_value(row):
    # 配对flag和对应的value
    paired = zip([row['A_flag'], row['B_flag'], row['C_flag']],
                 [row['A_value'], row['B_value'], row['C_value']])
    # 筛选flag为1的value
    selected = [val for flag, val in paired if flag == 1]
    # 按规则返回结果
    if len(selected) == 3:
        return pd.Series(selected).median()
    elif len(selected) == 2:
        return sum(selected) / 2
    else:
        return selected[0]

df['Final'] = df.apply(get_final_value, axis=1)

运行后即可得到与df1一致的Final列结果。

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