[Python/Pandas] 리트코드 1731. The Number of Employees Which Report to Each Employee

문제

Table: Employees
+-------------+----------+
| Column Name | Type     |
+-------------+----------+
| employee_id | int      |
| name        | varchar  |
| reports_to  | int      |
| age         | int      |
+-------------+----------+

employee_id is the column with unique values for this table.
This table contains information about the employees and the id of the manager they report to. Some employees do not report to anyone (reports_to is null). 
 
For this problem, we will consider a manager an employee who has at least 1 other employee reporting to them.
Write a solution to report the ids and the names of all managers, the number of employees who report directly to them, and the average age of the reports rounded to the nearest integer.
Return the result table ordered by employee_id.

 

https://leetcode.com/problems/the-number-of-employees-which-report-to-each-employee/description/?lang=pythondata

예시

Employees table:
+-------------+---------+------------+-----+
| employee_id | name    | reports_to | age |
+-------------+---------+------------+-----+
| 9           | Hercy   | null       | 43  |
| 6           | Alice   | 9          | 41  |
| 4           | Bob     | 9          | 36  |
| 2           | Winston | null       | 37  |
+-------------+---------+------------+-----+
Output: 
+-------------+-------+---------------+-------------+
| employee_id | name  | reports_count | average_age |
+-------------+-------+---------------+-------------+
| 9           | Hercy | 2             | 39          |
+-------------+-------+---------------+-------------+
Explanation: Hercy has 2 people report directly to him, Alice and Bob. Their average age is (41+36)/2 = 38.5, which is 39 after rounding it to the nearest integer.

Input: 
Employees table:
+-------------+---------+------------+-----+ 
| employee_id | name    | reports_to | age |
|-------------|---------|------------|-----|
| 1           | Michael | null       | 45  |
| 2           | Alice   | 1          | 38  |
| 3           | Bob     | 1          | 42  |
| 4           | Charlie | 2          | 34  |
| 5           | David   | 2          | 40  |
| 6           | Eve     | 3          | 37  |
| 7           | Frank   | null       | 50  |
| 8           | Grace   | null       | 48  |
+-------------+---------+------------+-----+ 
Output: 
+-------------+---------+---------------+-------------+
| employee_id | name    | reports_count | average_age |
| ----------- | ------- | ------------- | ----------- |
| 1           | Michael | 2             | 40          |
| 2           | Alice   | 2             | 37          |
| 3           | Bob     | 1             | 37          |
+-------------+---------+---------------+-------------+

문제 풀이

import pandas as pd

def count_employees(employees: pd.DataFrame) -> pd.DataFrame:
    # 'reports_to' 열 기준으로 그룹화하여 각 상사 아래의 직원 수와 평균 나이 계산
    # 'reports_count': 각 상사 아래 보고하는 직원 수 (size)
    # 'average_age': 각 상사 아래 직원들의 평균 나이 (mean)
    grouped = employees.groupby('reports_to').agg(reports_count=('employee_id', 'size'), average_age=('age', 'mean')).reset_index()
    
    # 'grouped' DataFrame을 원본 'employees' DataFrame과 병합
    merged = pd.merge(grouped, employees, how='inner', left_on='reports_to', right_on='employee_id')
    
    # 'merged' DataFrame에서 필요한 열만 선택하고, 'employee_id'를 기준으로 정렬
    result = merged[['employee_id', 'name', 'reports_count', 'average_age']].sort_values(by='employee_id')
    
    # 평균 나이 열의 값을 반올림 처리
    result['average_age'] = (result['average_age'] + 1e-9).round().astype(int)
    return result

---

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

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