如何用Pandas将多个虚拟变量合并为单列并计算类型均值？

下面提供两种简洁的pandas操作方法，帮你将虚拟变量合并为分类列，并计算各职位的平均薪资：

方法一：利用idxmax快速生成分类列 #

因为你的数据中每行仅有一个虚拟变量值为1，其余为0，所以可以直接取每行最大值对应的列名作为职位类型：

import pandas as pd
data = {'salary' : [50000, 45000, 55000, 40000, 35000, 45000, 30000,25000,35000],
        'manager': [1,1,1,0,0,0,0,0,0], 'foreman': [0,0,0,1,1,1,0,0,0], 
        'worker': [0,0,0,0,0,0,1,1,1]}
df = pd.DataFrame(data=data)

# 生成type列：提取每行值为1的虚拟变量列名
df['type'] = df[['manager', 'foreman', 'worker']].idxmax(axis=1)

# 保留需要的salary和type列
df_clean = df[['salary', 'type']]
print(df_clean)

输出的df_clean就是仅包含薪资和职位类型的数据集：

salary      type
0   50000   manager
1   45000   manager
2   55000   manager
3   40000   foreman
4   35000   foreman
5   45000   foreman
6   30000    worker
7   25000    worker
8   35000    worker

方法二：用melt重塑数据（适用于更复杂的虚拟变量场景） #

如果后续遇到多虚拟变量为1的情况，melt方法更通用，只需过滤值为1的行即可：

# 重塑数据，将虚拟变量列转为行
df_melt = df.melt(
    id_vars='salary',
    value_vars=['manager', 'foreman', 'worker'],
    var_name='type',
    value_name='is_type'
)

# 筛选出值为1的行，去掉无用列
df_clean = df_melt[df_melt['is_type'] == 1].drop('is_type', axis=1)

计算各职位平均薪资 #

得到df_clean后，直接用groupby计算平均：

avg_salary = df_clean.groupby('type')['salary'].mean().reset_index()
print(avg_salary)

输出结果：

type   salary
0   foreman  40000.0
1   manager  50000.0
2    worker  30000.0

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