Python中计算部门内员工与主管的4类薪资指标需求

核心思路 #

利用Pandas的groupby+transform/apply实现按部门分组，提取每个部门主管（serial_number='1'）的薪资数据后，基于此计算各类指标。

先做数据预处理，提取所有薪资列：

import pandas as pd

# 原始数据
df2=pd.DataFrame({
    'head_department': [ 'Anna','Anna', 'Anna','Anna','Anna',
                         'John', 'John','John',
                         'Denis', 'Denis',
                         'Sarah',
                         'Greg', 'Greg','Greg','Greg', 'Greg','Greg'],
    'serial_number': [1,2,3,4,5,1,2,3,1,2,1,1,2,3,4,5,6],
    'Worker_in_dpt': ['Anna','Anna-s_worker_2','Anna-s_worker_3','Anna-s_worker_4','Anna-s_worker_5',
                      'John','John-s worker_2','John-s_worker_3',
                      'Denis','Denis-s_worker_2',
                      'Sarah',
                      'Greg','Greg-s_worker_2', 'Greg-s_worker_3', 'Greg-s_worker_4','Greg-s_worker_5', 'Greg-s_worker_6'],
   'jan_salary': [1000, 600, 600, 500, 900, None, 600, 500, 1200, 800, 1400, None, 700, 600, 600, 450, 700],
   'feb_salary': [1100, 700, 700, 700, 800, None, 500, 400, 1800, 900, 1000, 1400, 900, 800, 500, 450, 700],
   'mar_salary': [1200, 800, 800, 900, 700, 1300, 600, 500, 1800, 600, 1100, 1600, 400, 700, 600, 250, 700],
   'apr_salary': [1600, 900, 900, 700, 700, 2300, 500, 400, 1800, 900, 1100, 1900, 200, 900, 500, 150, 700],
   'may_salary': [1100, 700, 700, 700, 800, 2300, 500, 400, 1800, 900, 1000, 1400, 900, 800, 500, 450, 700],
   'jun_salary': [1200, 800, 800, 900, 700, 1300, 500, 400, 1800, 900, 1000, 1400, 900, 800, 500, 450, 800],
   'jul_salary': [1000, 600, 600, 500, 900, 1300, 600, 500, 1200, 800, 1400, 1200, 700, 600, 600, 450, 700],
   'aug_salary': [1100, 700, 700, 700, 800, 2300, 600, 500, 1800, 600, 1100, 1600, 400, 700, 600, 250, None]
})

df2['serial_number'] = df2['serial_number'].astype('str')
df2['mean_salary_per_period'] = df2.mean(numeric_only=True, axis=1)

# 提取所有薪资月份列
salary_cols = [col for col in df2.columns if '_salary' in col]
# 提取各部门主管的薪资数据（serial_number='1'的行）
head_salaries = df2[df2['serial_number'] == '1'].set_index('head_department')[salary_cols]

1. 计算员工与部门主管的薪资相关性（corr_with_head_in_dpt） #

按部门分组，对每个员工的全周期薪资与对应主管的薪资做皮尔逊相关系数计算：

def calc_corr_with_head(group):
    # 获取当前部门主管的薪资数据
    head_salary_series = head_salaries.loc[group.name]
    # 对组内每个员工的薪资列，计算与主管薪资的相关系数
    return group[salary_cols].apply(lambda row: row.corr(head_salary_series), axis=1)

df2['corr_with_head_in_dpt'] = df2.groupby('head_department').apply(calc_corr_with_head).reset_index(drop=True)

2. 计算主管薪资为NaN期间的员工平均薪资（mean_when_1_in_NAN） #

先定位部门主管薪资为NaN的月份，仅计算员工在这些月份的薪资平均值；若主管无NaN值，则返回全周期平均：

def mean_when_head_nan(group):
    head_nan_mask = head_salaries.loc[group.name].isna()
    if head_nan_mask.any():
        # 筛选主管薪资为NaN的月份列
        nan_period_cols = salary_cols[head_nan_mask]
        return group[nan_period_cols].mean(axis=1)
    else:
        # 主管无NaN时返回全周期平均
        return group[salary_cols].mean(axis=1)

df2['mean_when_1_in_NAN'] = df2.groupby('head_department').apply(mean_when_head_nan).reset_index(drop=True)

3. 计算主管首次有薪资后的员工平均薪资（mean_when_1_not_in_NAN） #

找到主管首次出现非NaN薪资的月份，从该月份开始计算员工的薪资平均值；若主管无NaN值，则返回全周期平均：

def mean_from_first_head_non_nan(group):
    head_salary_series = head_salaries.loc[group.name]
    # 获取主管首次非NaN薪资的月份
    first_valid_month = head_salary_series.first_valid_index()
    if first_valid_month:
        # 获取从首次有效月份开始的所有薪资列
        start_idx = salary_cols.index(first_valid_month)
        target_cols = salary_cols[start_idx:]
        return group[target_cols].mean(axis=1)
    else:
        # 极端情况：主管全为NaN时返回全周期平均
        return group[salary_cols].mean(axis=1)

df2['mean_when_1_not_in_NAN'] = df2.groupby('head_department').apply(mean_from_first_head_non_nan).reset_index(drop=True)

4. 计算部门1月薪资总和（sum_all_dpt_in_jan） #

按部门分组，对jan_salary列求和（NaN值自动不参与求和，符合“包含主管薪资即使为NaN”的要求）：

df2['sum_all_dpt_in_jan'] = df2.groupby('head_department')['jan_salary'].transform('sum')

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