LeetCode 176：第二高薪资Pandas代码输出符合预期却报错求助

问题代码 #

import pandas as pd
def second_highest_salary(employee: pd.DataFrame) -> pd.DataFrame:
    employee.salary = employee.salary.drop_duplicates().reset_index(drop = True)
    employee = employee.sort_values(by = 'salary',ascending = False)
    if employee.shape[0] < 2:
        return pd.DataFrame(data = {'SecondHighestSalary':[None]})
    else:
        return pd.DataFrame(data = {'SecondHighestSalary':[employee.salary.iloc[[1]]]})

错误原因分析 #

• 错误1：修改原DataFrame列导致数据异常
  执行employee.salary = employee.salary.drop_duplicates().reset_index(drop=True)时，去重后的薪资列表长度可能小于原表行数，赋值回原列会让多余行的薪资变成NaN。比如原表有5行但只有2种不同薪资，去重后只剩2个值，原表后3行的salary会被设为NaN，此时employee.shape[0]还是5，你的判断条件employee.shape[0] < 2永远不会触发，即使实际只有1种薪资也会错误地去取第二高值。

• 错误2：返回值类型错误
  employee.salary.iloc[[1]]返回的是一个Series对象，而不是单个数值。这样构造的DataFrame里存储的是Series而非预期的薪资数值，LeetCode的判题器会识别为格式错误。应该用employee.salary.iloc[1]（不带方括号）来获取单个值。

修正后的代码示例 #

import pandas as pd
def second_highest_salary(employee: pd.DataFrame) -> pd.DataFrame:
    # 提取去重后的薪资并降序排序
    unique_salaries = employee['salary'].drop_duplicates().sort_values(ascending=False)
    # 判断是否有至少2种不同薪资
    if len(unique_salaries) < 2:
        return pd.DataFrame({'SecondHighestSalary': [None]})
    else:
        return pd.DataFrame({'SecondHighestSalary': [unique_salaries.iloc[1]]})

内容的提问来源于stack exchange，提问作者Yash Kumar Roy