본문 바로가기
Data Analysis/Query

리트코드 : 1789. Primary Department for Each Employee

by 베짱이28호 2024. 9. 7.

리트코드 : 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.

The result format is in the following example.



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             |
+-------------+---------------+
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.
  • 각 사람들의 메인 부서 찾기
  • 부서가 여러개면 Y FLAG를 찾아준다.
  • 아닌 경우 GROUP BY + COUNT에서 COUNT가 1인거를 UNION

문제 풀이

MySQL

SELECT EMPLOYEE_ID, DEPARTMENT_ID
FROM EMPLOYEE
GROUP BY EMPLOYEE_ID
HAVING COUNT(*) = 1

UNION

SELECT EMPLOYEE_ID, DEPARTMENT_ID
FROM EMPLOYEE
WHERE PRIMARY_FLAG = 'Y'
  • GROUP BY + HAVING으로 부서가 한 개인 사람만 구하기
  • WHERE 조건으로 PRIMARY 부서가 지정된 (부서 두 개 이상인) 사람 구하기

Pandas

import pandas as pd

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    employee.sort_values(by=['employee_id', 'primary_flag'], inplace=True)
    employee.drop_duplicates(subset='employee_id', keep='last', inplace=True)
    return employee[['employee_id', 'department_id']]
  • sort_values + drop_duplicates
  • y는 각 사람별로 한 번만 나온다는점을 생각해서 사람별로 flag를 정렬한다.
  • keep last를 걸어놓으면 마지막이 y인 부서를 가져올 수 있다.

코멘트

  • .

댓글