Problem

Table: Employee

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

+---------------+---------+
| 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.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

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.

Solution

Method 1 - Primary Selection (Flag or Single Department)

Intuition

We need to find the primary department for each employee. If an employee has only one department, that department is primary. If they have multiple, the one with primary_flag = 'Y' is primary; use employee_id and department_id to identify rows.

Approach

  1. For each employee_id, if they have only one department_id, select that department.
  2. If they have multiple, select the row where primary_flag = ‘Y’.
  3. Combine both cases in a single query using SQL or pandas.

Complexity

  • ⏰ Time complexity: O(N) where N is the number of rows in Employee.
  • 🧺 Space complexity: O(N) for intermediate grouping.

Code

1
2
3
4
5
6
7
8
9

SELECT employee_id, department_id
FROM Employee
WHERE primary_flag = 'Y'
   OR employee_id IN (
        SELECT employee_id
        FROM Employee
        GROUP BY employee_id
        HAVING COUNT(*) = 1
   );
1
2
3
4
5
6
7
8
9

SELECT employee_id, department_id
FROM Employee
WHERE primary_flag = 'Y'
   OR employee_id IN (
        SELECT employee_id
        FROM Employee
        GROUP BY employee_id
        HAVING COUNT(*) = 1
   );
 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import pandas as pd

# Assume df is the Employee DataFrame
def primary_department(df: pd.DataFrame) -> pd.DataFrame:
    # Employees with only one department
    single = df.groupby('employee_id').filter(lambda x: len(x) == 1)
    # Employees with multiple departments, pick primary_flag == 'Y'
    multi = df[df['primary_flag'] == 'Y']
    result = pd.concat([single, multi])
    return result[['employee_id', 'department_id']]