Leetcode 解鎖 : Primary Department for Each Employee

題目概述

問題連結

資料中，員工可能同時隸屬於多個部門，但只有一個主要部門（標記為 primary_flag = 'Y'）。如果員工只屬於一個部門，那麼這個部門就自動成為主要部門，不論 primary_flag 的值是什麼。

這就像是在處理公司內部的人事資料時，需要確定每位員工的主要歸屬部門，以便進行績效評估、資源分配等管理工作。

Input 格式

表格名稱為 Employee，包含以下欄位：

• employee_id - 員工ID
• department_id - 部門ID
• primary_flag - 是否為主要部門（'Y' 或 'N'）

每位員工可能有多筆記錄（代表隸屬多個部門），但最多只有一個部門會被標記為主要部門（primary_flag = 'Y'）。

Output 格式

返回每位員工的主要部門，包含兩個欄位：

• employee_id - 員工ID
• department_id - 主要部門ID

Pandas 解法 1

排序後取第一筆

Step 1: 排序資料，讓 primary_flag='Y' 的記錄排在前面

Step 2: 對每個 employee_id 只保留第一筆記錄

Step 3: 只返回需要的欄位

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
   
    df_sorted = employee.sort_values(
        by=['employee_id', 'primary_flag'], 
        ascending=[True, False]
    )
    
    result = df_sorted.drop_duplicates(
        subset='employee_id', 
        keep='first'
    )
    
    return result[['employee_id', 'department_id']]

這個方法先將資料依照 employee_id 和 primary_flag 排序，確保主要部門（'Y'）排在前面，然後用 drop_duplicates 只保留每位員工的第一筆記錄。

Pandas 解法 2

計數 + 條件篩選

Step 1: 計算每位員工的部門數量

Step 2: 篩選條件：只有一個部門 或 標記為主要部門

Step 3: 應用篩選條件並只返回需要的欄位

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:

    employee['dept_count'] = employee.groupby('employee_id')['department_id'].transform('count')
    filt = (employee['dept_count'] == 1) | (employee['primary_flag'] == 'Y')
    result = employee.loc[filt, ['employee_id', 'department_id']]
    
    return result

這個方法先計算每位員工有幾個部門，然後篩選出「只有一個部門」或「標記為主要部門」的記錄。

Pandas 解法 3

巧用 duplicated 函數

Step 1: 沒有重複的員工ID 或 標記為主要部門

Step 2: 應用篩選條件並只返回需要的欄位

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    
    filt = (~employee.duplicated(subset='employee_id', keep=False)) | (employee['primary_flag'] == 'Y')

    result = employee.loc[filt, ['employee_id', 'department_id']]
    
    return result

Pandas 解法 4

分開處理後合併

Step 1: 找出所有標記為主要部門的記錄

Step 2: 找出只有一個部門的員工記錄（不在主要部門列表中的員工）

Step 3: 合併兩部分結果

Step 4: 只返回需要的欄位

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    
    primary_depts = employee.loc[employee['primary_flag'] == 'Y', :].copy()
    single_dept_employees = employee[~employee['employee_id'].isin(primary_depts['employee_id'])]
    result = pd.concat([primary_depts, single_dept_employees], ignore_index=True)

    return result[['employee_id', 'department_id']]

Pandas 解法 5

使用 assign + query 的函數式寫法

使用 assign 生成新欄位，並用 query 篩選

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    
    result = employee.assign(
        max_flag = lambda df: df.groupby('employee_id')['primary_flag']
                             .transform('max')
    ) .query('primary_flag == max_flag') [['employee_id', 'department_id']]
    
    return result

Pandas 解法 6

使用 merge 合併計數結果

Step 1: 計算每位員工的部門數量

Step 2: 合併計數結果與原始資料

Step 3: 將只有一個部門的員工的 primary_flag 設為 'Y'

Step 4: 篩選出標記為主要部門的記錄

def find_primary_department(employee: pd.DataFrame) -> pd.DataFrame:
    # 
    dept_counts = employee.groupby('employee_id').count()[['department_id']].reset_index()
    dept_counts = dept_counts.rename(columns={'department_id': 'dept_count'})
    merged_data = pd.merge(dept_counts, employee, on='employee_id', how='inner')
    merged_data.loc[merged_data['dept_count'] == 1, 'primary_flag'] = 'Y'
    result = merged_data[merged_data['primary_flag'] == 'Y'][['employee_id', 'department_id']]
    
    return result

謝謝您花時間將此篇文章讀完，若覺得對您有幫助可以幫忙按個讚、分享來或是珍藏喔!也歡迎Follow我的Threads/ FB，持續追蹤生產力工具、商業分析、商業英文的實用範例，提升自己的職場力喔!