Python/알고리즘

[파이썬 알고리즘 인터뷰] 31. 상위 K 빈도 요소

AI 꿈나무 2021. 3. 5. 15:28
반응형

 

 이 포스팅은 파이썬 알고리즘 인터뷰를 공부하면서 정리한 내용입니다.

 

 출처 : 파이썬 알고리즘 인터뷰

 코드 출처 : 파이썬 알고리즘 인터뷰 깃허브

 


31. 상위 K 빈도 요소

leetcode 346. Top K Frequent Elements 문제입니다.

leetcode.com/problems/top-k-frequent-elements/

 

 

풀이

1. 우선순위 큐 이용

 파이썬에서 우선순위 큐는 heapq 모듈을 사용합니다. nums의 빈도수를 Counter 모듈로 저장한뒤에, heapq에 빈도수를 음수로 저장합니다. heapq는 최소합을 제공하므로 음수로 저장해야 합니다. k번 반복하여 heapq를 pop하면 빈도수가 많은 요소순대로 추출이 됩니다.

 

def topKFrequent(self, nums, k):
    freqs = collections.Counter(nums)
    heap = []

    for f in freqs:
        heapq.heappush(heap, (-freqs[f], f))

    result = []
    for _ in range(k):
        result.append(heapq.heappop(heap)[1])

    return result

 

2. 한줄 풀이

 Counter의 .most_common 함수를 사용하는 풀이입니다. 이 함수를 사용하면 요소와 빈도수를 반환하는데, zip 함수를 이용하여 요소끼리만 묶어 반환합니다.

 

def topKFrequent(self, nums, k):
    return list(zip(*collections.Counter(nums).most_common(k)))[0]
반응형