설명

LeetCode의 Top K Frequent Elements 문제다.

숫자들로 이루어진 배열과 정수 k가 주어질 때 배열에서 가장 많이 등장한 정수를 내림차순으로 k개 반환하는 것이 목적이다. 헷갈릴 수 있지만 배열에 k개 이상 들어있는 정수를 반환하는 게 아니라 배열에 들어있는 정수들 중에서 많이 들어있는 순서대로 k개를 반환하라는 것이다.

풀이

사전 자료형 활용

import collections


class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        counter = collections.defaultdict(int)
        for num in nums:
            counter[num] += 1
            
        result = sorted(counter, key=lambda k:counter[k])
        
        return result[-k:]

설명에서 헤매지 않았다면 어렵지 않게 풀 수 있다. sorted 메서드를 활용하여 사전의 키-값 쌍에서 값, 즉 빈도로 오름차순 정렬하고 마지막 k개를 반환한다.

우선순위 큐(힙) 활용

파이썬의 heapq 모듈을 활용하여 우선순위 큐에 삽입한 후 k개 만큼 빼서 반환하는 방식이다.

import heapq
import collections


class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        counter = collections.defaultdict(int)
        for num in nums:
            counter[num] += 1
        
        heap = []
        for num in counter:
            heapq.heappush(heap, (-counter[num], num))
            
        result = []    
        for _ in range(k):
            result.append(heapq.heappop(heap)[1])
            
        return result

사전 자료형으로 정수의 빈도를 확인한 후 튜플에 빈도를 음수로 변환하여 힙에 삽입하고 있다. 이전에도 언급했지만 heapq 모듈은 자료구조를 오름차순으로 정렬한다. 그렇기 때문에 내림차순으로 정렬하려면 저렇게 튜플로 감싸서 반대로 정렬될 수 있도록 조작할 수 있다.

예를 들어 숫자 1이 3개, 숫자 2가 5개, 숫자 3이 1개 들어있다고 할 때 그냥 튜플로 감싼다면 (3, 1), (5, 2), (1, 3)이며 튜플의 첫번째 원소로 비교해서 (1, 3), (3, 1), (5, 2) 순서대로 우선순위 큐에 저장되기 때문에 높은 빈도대로 k개의 원소를 꺼내기 어렵다.

그래서 "-counter[num]"처럼 빈도를 음수로 변환해서 (-3, 1), (-5, 2), (-1, 3)로 삽입한다면 (-5, 2), (-3, 1), (-1, 3)으로 정렬되기 때문에 간단하게 k번 꺼내서 높은 빈도순으로 반환할 수 있다.