[LeetCode - Pandas] (Easy) 1789. Primary Department for Each Employee

Table: Employee

+---------------+---------+
| Column Name   |  Type   |
+---------------+---------+
| employee_id   | int     |
| department_id | int     |
| primary_flag  | varchar |
+---------------+---------+

(employee_id, department_id)는 이 테이블의 기본 키(고유 값이 있는 열의 조합)입니다.
employee_id는 직원의 ID입니다.
department_id는 직원이 속한 부서의 ID입니다.
primary_flag는 'Y' 또는 'N' 값으로 이루어진 ENUM(카테고리)입니다. 만약 플래그가 'Y'라면, 해당 부서는 직원의 주요 부서입니다. 플래그가 'N'이라면, 해당 부서는 직원의 주요 부서가 아닙니다.

 

문제

직원은 여러 부서에 속할 수 있습니다. 직원이 다른 부서에 가입할 때, 어떤 부서를 주요 부서로 선택할지 결정해야 합니다. 단, 직원이 하나의 부서에만 속한 경우, 그들의 주요 부서 열(primary column)은 'N'입니다.

모든 직원과 그들의 주요 부서를 보고하는 해결책을 작성하세요. 하나의 부서에만 속한 직원은 그들의 유일한 부서를 보고하세요.

결과 테이블은 예시와 같은 형식으로 반환됩니다.

Example 1:

Input: 
Employee table:
+-------------+---------------+--------------+
| employee_id | department_id | primary_flag |
+-------------+---------------+--------------+
| 1           | 1             | N            |
| 2           | 1             | Y            |
| 2           | 2             | N            |
| 3           | 3             | N            |
| 4           | 2             | N            |
| 4           | 3             | Y            |
| 4           | 4             | N            |
+-------------+---------------+--------------+
Output: 
+-------------+---------------+
| employee_id | department_id |
+-------------+---------------+
| 1           | 1             |
| 2           | 1             |
| 3           | 3             |
| 4           | 3             |
+-------------+---------------+

설명 :

- 직원 1의 주요 부서는 1번 부서입니다.

- 직원 2의 주요 부서는 1번 부서입니다.

- 직원 3의 주요 부서는 3번 부서입니다.

- 직원 4의 주요 부서는 3번 부서입니다.

 

✏️ 풀이

import pandas as pd

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    employee['cnt'] = employee.groupby('employee_id')['employee_id'].transform('count')

    return employee[(employee['cnt'] == 1) | ((employee['cnt'] > 1) & (employee['primary_flag'] == 'Y'))][['employee_id', 'department_id']]

 

📚 개념정리

transform() ? 데이터 프레임의 각 열에 여러 함수를 동시에 적용하고, 결과를 동일한 행의 데이터로 반환하는 메서드

transform을 사용하지 않느다면 위의 employee 데이터 프레임에서 employee_id별 count를 구하고 그걸 다시 각 employee_id에 merge를 해줘야하는 번거로운 과정을 거쳐야하는데 transform을 사용하면 employee_id별 count 값을 해당 데이터 프레임의 동일한 행에 데이터를 위치 시킴으로 써 지곤 데이터 프레임과 동일한 사이즈의 값을 얻게 된다

 

참고 블로그 : https://m.blog.naver.com/sw4r/222392753166