[Python/Pandas] 리트코드 1789. Primary Department for Each Employee

문제

Table: Employee
+---------------+---------+
| Column Name   |  Type   |
+---------------+---------+
| employee_id   | int     |
| department_id | int     |
| primary_flag  | varchar |
+---------------+---------+
(employee_id, department_id) is the primary key (combination of columns with unique values) for this table.
employee_id is the id of the employee.
department_id is the id of the department to which the employee belongs.
primary_flag is an ENUM (category) of type ('Y', 'N'). If the flag is 'Y', the department is the primary department for the employee. If the flag is 'N', the department is not the primary.

Employees can belong to multiple departments. When the employee joins other departments, they need to decide which department is their primary department. Note that when an employee belongs to only one department, their primary column is 'N'.

Write a solution to report all the employees with their primary department. For employees who belong to one department, report their only department.

Return the result table in any order.

 

https://leetcode.com/problems/primary-department-for-each-employee/description/?lang=pythondata

예시

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             |
+-------------+---------------+
Explanation: 
- The Primary department for employee 1 is 1.
- The Primary department for employee 2 is 1.
- The Primary department for employee 3 is 3.
- The Primary department for employee 4 is 3.

문제 풀이

import pandas as pd

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    # 각 직원이 속한 부서 수를 계산하여 'cnt_department_id' 열 생성
    employee['cnt_department_id'] = employee.groupby('employee_id')['department_id'].transform('count')
    
    # 조건을 만족하는 직원의 'employee_id'와 'department_id'를 필터링
    result = employee.loc[(employee['primary_flag'] == 'Y') | (employee['cnt_department_id'] == 1), ['employee_id', 'department_id']]
    return result

---

파이썬을 독학하시는 분들에게 도움이 되길 바라며,

혹 더 좋은 방법이 있거나 오류가 있다면 편하게 말씀 부탁드립니다.