El número de empleados que reportan a cada empleado

Para este problema, consideraremos que un empleado es gerente si tiene al menos 1 empleado a su cargo. Escribe una solución para reportar los IDs y los nombres de todos los gerentes, el número de empleados que tienen a su cargo, y la edad promedio de los empleados subordinados redondeada al valor entero más cercano. Regresa el resultado ordenado por employee_id.

Tabla: Employees

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

- employee_id es una columna con valores únicos is the column with unique values for this table.
- Esta tabla contiene información de los empleados y el ID del gerente al que reportan. Algunos empleados no
  reportan a nadie (reports_to es null).

El formato del resultado se muestra en el siguiente ejemplo.

Ejemplo 1:

Entrada:

Tabla Employees:
+-------------+---------+------------+-----+
| employee_id | name    | reports_to | age |
+-------------+---------+------------+-----+
| 9           | Hercy   | null       | 43  |
| 6           | Alice   | 9          | 41  |
| 4           | Bob     | 9          | 36  |
| 2           | Winston | null       | 37  |
+-------------+---------+------------+-----+

Salida:

+-------------+-------+---------------+-------------+
| employee_id | name  | reports_count | average_age |
+-------------+-------+---------------+-------------+
| 9           | Hercy | 2             | 39          |
+-------------+-------+---------------+-------------+

Explicación: Hercy tiene dos personas a su cargo, Alice y Bob. Su promedio de edad es (41+36)/2 = 38.5, el
cual se convierte en 39 después de redondearlo al valor entero más próximo.

Solución

import pandas as pd


def count_employees(employees: pd.DataFrame) -> pd.DataFrame:
    managers_data = employees.groupby('reports_to', as_index=False).agg(
        average_age=('age', 'mean'),
        reports_count=('reports_to', 'size')
    )

    manager_names = employees[
        employees['employee_id'].isin(managers_data['reports_to'])
    ]

    merged = pd.merge(managers_data,
                      manager_names,
                      left_on='reports_to',
                      right_on='employee_id')

    result = pd.DataFrame({
        'employee_id': merged['reports_to_x'],
        'name': merged['name'],
        'reports_count': merged['reports_count'],
        'average_age': round(merged['average_age'] + 1e-6).astype(int),
    })

    result.sort_values(by='employee_id', inplace=True)

    return result


def test_count_employees():
    columns = {
        'employee_id': 'Int64',
        'name': 'object',
        'reports_to': 'Int64',
        'age': 'Int64'
    }

    data = [
        [9, 'Hercy', None, 43],
        [6, 'Alice', 9, 41],
        [4, 'Bob', 9, 36],
        [2, 'Winston', None, 37]
    ]

    employees = pd.DataFrame(data,
                             columns=columns.keys()).astype(columns)

    got = count_employees(employees)
    expected = pd.DataFrame({
        'employee_id': [9],
        'name': ['Hercy'],
        'reports_count': [2],
        'average_age': [39],
    })

    pd.testing.assert_frame_equal(got, expected, check_dtype=False)