[리트코드/leetcode/python] 347. Top K Frequent Elements

이미지를 누르면 문제 링크를 보실 수 있습니다

이 문제를 처음 봤을 때 계수정렬을 생각해봤습니다.

nums[i]의 범위가 -10^4 부터 시작되므로, 음수도 포함될 수 도 있기에, 음수일 때와 양수일 때를 나눠서 계수정렬로 로직을 짰었습니다. 하지만, 계수정렬을 한 리스트가 내림차순되면서, k 번째까지 리스트를 슬라이싱 하는데에, 원소 순서가 변하기에 한계가 있었습니다.

 

어떤 수가 몇 번 나왔는지를 체크해야 하므로, 계수정렬보다는 딕셔너리 등을 사용하는 것이 더 적합했습니다.

follow up을 보면, O(n log n)의 시간복잡도를 구현해보라고 하는데, 딕셔너리로 풀 때, O(n)으로 풀리는 것이 아닌가 하는 의구심으로 문제를 풀기 시작했습니다.

 

푼 코드는 아래와 같습니다.

class Solution(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """

        dicts = defaultdict(int)
        
        for i in nums:
            dicts[i] += 1

        dicts = sorted(dicts.items(),key = lambda x : -x[1])
       
        answer = []
        for j in range(k):
            answer.append(dicts[j][0])

        return answer

 

defaultdict(int)를 씀으로써, 초기 int형 딕셔너리를 생성하고, nums 를 순회했습니다.

nums를 순회하면서, dicts[i] 에 해당하는 횟수를 1씩 증가하면서, 해당 nums에 어떤 수가 몇 번 나왔는지를 체크했습니다.

그리고, values()를 기준으로 내림차순을 하고, 상위 k 번째에 해당하는 key값을 뽑아내고 리스트에 담으면 그것이 답이 되는 로직입니다.

 

extends()는 sort()처럼 반환값이 없는 내장함수인데, 이번 문제를 풀면서 그 사실에 대해 확인했고, key = lambda x : x[1] 등의 표현법도 values를 기준으로 정렬할 때 사용하는 것임을 재차 확인했습니다.

 

heapq 라이브러리를 쓰는 다른 코드도 한 번 살펴보겠습니다. minheap에서 nlargest() 함수를 사용한 것이 특징입니다.

 

이상으로 Top K Frequent Elements 문제를 정리해봤습니다.