[Python/Pandas] 그룹별 상위 N개 데이터 추출하기 (groupby, cumcount)

안녕하세요! 이번에는 Pandas를 사용하여 특정 그룹별 상위 N개의 데이터를 추출하는 방법을 소개합니다. 데이터 분석 작업 중에 그룹화된 데이터에서 상위 몇 개의 행만 선택하고자 하는 경우가 자주 발생합니다. 이번 글에서는 두 가지 방법으로 그룹별 상위 데이터를 추출해보겠습니다.

 

---

 

예제 데이터

이번에 비트코인 가격이 많이 올라왔던 게 생각나서 데이터는 캐글에 있는 코인 가격에 대한 데이터 중 일부를 활용해보았습니다. (출처: Kaggle, Top 100 Cryptocurrencies Historical Dataset)

• Name: 코인의 이름 (Bitcoin, Ethereum...)
• Date: 가격 데이터의 날짜
• Close: 해당 날짜의 종가
• 데이터 행수: 13,474

import os
import pandas as pd

# 코인 가격 데이터 가져오기
folder_url = './Coin dataset/'  # 데이터가 저장된 폴더 경로
file_list = os.listdir(folder_url)  # 폴더 내 파일 리스트 가져오기
df_list = list()

# 각 파일을 읽어서 병합
for file_name in file_list:
    df_coin = pd.read_csv(folder_url + file_name)  # 파일 읽기
    df_coin['Name'] = file_name.split('.')[0]     # 코인 이름 추가
    df_list.append(df_coin)

# 모든 데이터를 하나의 DataFrame으로 병합
df = pd.concat(df_list, ignore_index=True)

# 데이터 확인
df.head()

 

그룹별 상위 2개 데이터 추출 방법

방법 1: groupby().head() 사용

groupby()와 head() 메서드를 사용하면 각 그룹의 상위 N개의 데이터를 간단히 추출할 수 있습니다.

result1 = df.groupby('Name').head(2)

 

방법 2: cumcount()를 사용한 SQL의 ROW_NUMBER()와 유사한 방식

SQL에서 사용하는 ROW_NUMBER() 함수와 비슷하게, groupby()와 cumcount() 메서드를 활용하여 그룹별로 순서를 부여할 수 있습니다. 이후 조건을 사용해 상위 N개의 데이터를 필터링하는 방식입니다.

df['rank'] = df.groupby('Name').cumcount() + 1
result2 = df[df['rank'] <= 2]

 

결과 비교

방법 2의 결과물의 경우 cumcount()를 통해 부여한 순서열이 추가되어 있습니다.

---

 

DataFrame에서 그룹별로 상위 N개의 데이터를 선택해야 하는 상황이 있어 알아보았던 방법을 소개드렸습니다. 저처럼 그룹별 데이터를 확인하고자 하는 분들에게 도움이 되길 바랍니다. 오늘도 읽어주셔서 감사합니다!