[Python/Pandas] 리트코드 1075. Project Employees I

문제

Table: Project
+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| project_id  | int     |
| employee_id | int     |
+-------------+---------+

(project_id, employee_id) is the primary key of this table.
employee_id is a foreign key to Employee table.
Each row of this table indicates that the employee with employee_id is working on the project with project_id.

Table: Employee
+------------------+---------+
| Column Name      | Type    |
+------------------+---------+
| employee_id      | int     |
| name             | varchar |
| experience_years | int     |
+------------------+---------+

employee_id is the primary key of this table. It's guaranteed that experience_years is not NULL.
Each row of this table contains information about one employee.
 
Write an SQL query that reports the average experience years of all the employees for each project, rounded to 2 digits.
Return the result table in any order.

 

https://leetcode.com/problems/project-employees-i/description/?lang=pythondata

예시

Input: 
Project table:
+-------------+-------------+
| project_id  | employee_id |
+-------------+-------------+
| 1           | 1           |
| 1           | 2           |
| 1           | 3           |
| 2           | 1           |
| 2           | 4           |
+-------------+-------------+
Employee table:
+-------------+--------+------------------+
| employee_id | name   | experience_years |
+-------------+--------+------------------+
| 1           | Khaled | 3                |
| 2           | Ali    | 2                |
| 3           | John   | 1                |
| 4           | Doe    | 2                |
+-------------+--------+------------------+
Output: 
+-------------+---------------+
| project_id  | average_years |
+-------------+---------------+
| 1           | 2.00          |
| 2           | 2.50          |
+-------------+---------------+

Explanation: The average experience years for the first project is (3 + 2 + 1) / 3 = 2.00 and for the second project is (3 + 2) / 2 = 2.50

문제 풀이

import pandas as pd

def project_employees_i(project: pd.DataFrame, employee: pd.DataFrame) -> pd.DataFrame:
	# employee_id 기준으로 merge
    merged = pd.merge(project, employee, how='left', on='employee_id')
    
    # project_id 별 experience_years 평균을 구하고 round
    grouped = merged.groupby('project_id')['experience_years'].apply(lambda x: round(x.mean(), 2)).reset_index()
    
    # 문제 요구 출력에 맞추어 이름 변경
    result = grouped.rename(columns={'experience_years': 'average_years'})
    
    return result

---

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

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