Python 알고리즘 인터뷰 - 문제 31. 상위 K 빈도 요소

leetcode 347. Top K Frequenct Elements - https://leetcode.com/problems/top-k-frequent-elements/description/

 

상위 k번 이상 등장하는 요소를 추출하라

 

예제1

입력

nums = [1, 1, 1, 2, 2, 3], k = 2

 

출력

[1, 2]

 

예제2

입력

nums = [1, 2], k = 2

 

출력

[1, 2]

 

풀이

Counter와 heap 정렬을 사용한 풀이

리스트에서 가장 '빈번하게' 존재하는 숫자 순서대로 k개를 추출하는 문제이다. 이를 위해서는 리스트의 숫자별 개수를 구하기 위해 Counter 딕셔너리를 사용한다. Counter는 다음과 같이 값이 저장된다.

 

>>> freqs = Counter(nums)

>>> print(freqs)

Counter({1: 3, 2: 2, 3: 1})

 

상위 k개를 추출하기 위해서는 Counter에 개수를 기준으로 정렬을 하여 상위 k를 추출하면 된다. 이를 위해 힙정렬을 사용하며, 힙은 기본적으로 작은 값이 먼저 추출되도록 구현되어 있으므로, 개수를 음수로 하여 힙에 저장하면 개수가 가장 큰 순서대로 저장된다.

추후에 key 값 추출을 위해 (-value, key) tuple로 heap에 push한다.

코드는 다음과 같다.

heap_freq = []
freqs = Counter(nums)

for key, value in freqs.items():
    heapq.heappush(heap_freq, (-value, key))

 

이후 k번 heap에 데이터를 pop한 후 tuple에 key를 topk 리스트에 저장한다.

코드는 다음과 같다.

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

return topk

 

전체 코드는 다음과 같다.

def topKFrequent(nums: list, k):
    heap_freq = []
    freqs = Counter(nums)

    for key, value in freqs.items():
        heapq.heappush(heap_freq, (-value, key))

    topk = []
    k = min(k, len(heap_freq))
    for _ in range(k):
        topk.append(heapq.heappop(heap_freq)[1])

    return topk

 

Pythonic한 풀이

Counter에서 most_common(k) 함수를 호출하면 상위 k개 만큼 개수가 많은 (key, value)쌍의 tuple 리스트를 출력한다.

예를 들어 nums = [1, 1, 1, 2, 2, 3], k = 2의 경우

>>> counter = Counter([1, 1, 1, 2, 2, 3]).most_common(2)

[(1, 3), (2, 2)]

으로 출력한다.

여기서 [(1, 3), (2, 2)] 만 묶어서 리스트로 만들어 출력하면 상위 k의 값 list를 출력할 수 있다. 이를 위해 zip함수를 사용하여 most_common함수의 리턴 튜플 리스트를 언패킹하여 리스트로 출력하면 다음과 같다.

>>> list(zip(*counter.most_common(2)))

[(1, 2), (3, 2)]

여기서 리스트의 첫번째 인덱스의 튜플만 리턴해주면 상위 k의 값을 구할 수 있다. 코드는 다음과 같다.

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