[Python/Pandas] 리트코드 1179. Reformat Department Table

문제

Table: Department
+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| id          | int     |
| revenue     | int     |
| month       | varchar |
+-------------+---------+

In SQL,(id, month) is the primary key of this table.
The table has information about the revenue of each department per month.
The month has values in ["Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"].
 
Reformat the table such that there is a department id column and a revenue column for each month.
Return the result table in any order.

 

https://leetcode.com/problems/reformat-department-table/description/?lang=pythondata

예시

Input: 
Department table:
+------+---------+-------+
| id   | revenue | month |
+------+---------+-------+
| 1    | 8000    | Jan   |
| 2    | 9000    | Jan   |
| 3    | 10000   | Feb   |
| 1    | 7000    | Feb   |
| 1    | 6000    | Mar   |
+------+---------+-------+
Output: 
+------+-------------+-------------+-------------+-----+-------------+
| id   | Jan_Revenue | Feb_Revenue | Mar_Revenue | ... | Dec_Revenue |
+------+-------------+-------------+-------------+-----+-------------+
| 1    | 8000        | 7000        | 6000        | ... | null        |
| 2    | 9000        | null        | null        | ... | null        |
| 3    | null        | 10000       | null        | ... | null        |
+------+-------------+-------------+-------------+-----+-------------+

Explanation: The revenue from Apr to Dec is null.
Note that the result table has 13 columns (1 for the department id + 12 for the months).

문제 풀이

import pandas as pd

def reformat_table(department: pd.DataFrame) -> pd.DataFrame:
    prefixes = ["Jan","Feb","Mar","Apr","May","Jun","Jul","Aug","Sep","Oct","Nov","Dec"]
    
    # 피봇 테이블 생성
    pivoted = department.pivot(index='id', columns='month', values='revenue')
    
    # 열 reindex
    result = pivoted.reindex(columns=prefixes)
    
    # 열 이름 변경
    result.rename(columns=lambda prefix: prefix + '_Revenue', inplace=True)
    
    # 인덱스를 열로 변환 (id열)
    result.reset_index(inplace=True)
    
    return result

---

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

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