[코딩테스트]/[Pandas]

[LeetCode - Pandas] (Medium) 177. Nth Highest Salary

잰잰' 2025. 5. 13. 19:57
Table: Employee

+-------------+------+
| Column Name | Type |
+-------------+------+
| id          | int  |
| salary      | int  |
+-------------+------+

id는 이 테이블의 기본 키(고유한 값을 갖는 열)입니다.
이 테이블의 각 행은 직원의 급여에 대한 정보를 포함하고 있습니다.

 

문제

Employee 테이블에서 n번째로 높은 서로 다른 급여를 찾는 솔루션을 작성하세요.
서로 다른 급여가 n개보다 적으면 null을 반환하세요.

결과 형식은 다음 예시와 같습니다.

Example 1:

Input: 
Employee table:
+----+--------+
| id | salary |
+----+--------+
| 1  | 100    |
| 2  | 200    |
| 3  | 300    |
+----+--------+
n = 2
Output: 
+------------------------+
| getNthHighestSalary(2) |
+------------------------+
| 200                    |
+------------------------+
Example 2:

Input: 
Employee table:
+----+--------+
| id | salary |
+----+--------+
| 1  | 100    |
+----+--------+
n = 2
Output: 
+------------------------+
| getNthHighestSalary(2) |
+------------------------+
| null                   |
+------------------------+

 

✏️ 풀이

import pandas as pd

def nth_highest_salary(employee: pd.DataFrame, N: int) -> pd.DataFrame:
    unique_salary = employee[['salary']].drop_duplicates()
    unique_salary['rnk'] = unique_salary['salary'].rank(ascending=False)

    new_col = f'getNthHighestSalary({N})'

    if len(unique_salary) < N :
        return pd.DataFrame({new_col: [None]})

    result = unique_salary[unique_salary['rnk'] == 2][['salary']].rename(columns={'salary' : new_col})

    return result

이전 176번과 풀이는 상당 부분 유사하다

1. salary 컬럼명을 getHthHighestSalary(N) 으로 변경해야 하므로 해당 포맷을 저장할 new_col을 만들어준다

2. 데이터프레임의 row수가 N보다 작으면 Null 리턴

3. 아닐 경우 순위가 2위인 데이터 리턴