[python3] 347. Top K Frequent Elements

풀이 일자: 23.05.22

난이도: [Medium]

분류: [Hash Table, Divide and Conquer, Sorting, Heap, Bucket Sort]

문제 내용

단순하게 숫자 배열이 주어졌을 때 상위 k개의 빈번하게 등장하는 수를 순서상관 없이 리턴하는 문제였다.

 

Follow-up에서는 NLogN의 시간복잡도를 갖는 알고리즘으로 풀어보라고 했다.

 

문제 해결 흐름

 

1. 단순하게 생각해보면 dictionary 하나 만들어서 숫자 카운트하고 이를 value기준으로 sort해서 return하면 끝나겠다.

 

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        cnt = dict();

        for i in range(len(nums)):
            if nums[i] not in cnt:
                cnt[nums[i]] = 0;
            cnt[nums[i]] += 1;
        
        cnt = sorted(cnt.items(), key = lambda item: item[1], reverse=True);

        return [cnt[i][0] for i in range(k)];

 

Time Complexity:  O(NLogN)

가장 시간에 영향을 많이 주는 알고리즘은 sorted()함수이다. 일반적으로 프로그래밍 언어에 내장되어있는 sort함수는 대부분 quick sort를 기반으로 구현이 되어있다고 알고있는데 파이썬의 sorted함수는 Insert sort와 Merge sort를 합친 Timsort라는 알고리즘을 사용한다. 더 자세한 내용에 대해 알고 싶다면 이 네이버 기술 블로그를 이용하면 된다. 이는 최선의 시간복잡도가 O(N)이고 평균과 최악의 시간복잡도는 O(NLogN)으로 효과적인 정렬알고리즘이다.

 

 

Space Complexity: O(N)

현재 숫자를 카운트하기 위해서 dictionary를 사용했기 때문에 배열 내의 Unique한 숫자의 갯수만큼이 공간복잡도가 될 것이다. 

 

 

2. 만약 다른 알고리즘을 사용한다면 NlogN의 시간복잡도를 가진 힙정렬이 떠오를 것이다. 

→ Quick 정렬은 pivot을 잘못 설정하면 N^2의 최악의 시간복잡도를 가지기도 하므로 힙정렬이 안전하다.

 

1) 우선 앞선 방법과 같이 cnt 딕셔너리를 이용하여 갯수를 센다

2) 이후 대상 숫자와 갯수의 자리를 바꾼다. 

=> 여기서 갯수에 -를 붙이는 이유는  python에서 heappush를 하게 되면 min heap을 기준으로 들어가기 때문에 max heap이 되려면 -를 붙여 큰 값이 가장 작게 만들어 주는 것이다.

3) 그래로 heappush를 한 후 k개를 꺼낸다.

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        cnt = dict();
        heap = [];
        ans = [];

        for i in range(len(nums)):
            if nums[i] not in cnt:
                cnt[nums[i]] = 0;
            cnt[nums[i]] += 1;
        
        cnt = [(-j , i) for i, j in cnt.items()]

        for item in cnt:
            heappush(heap, item);

        for _ in range(k):
            ans.append(heappop(heap)[1])

        return ans;

Time Complexity:  O(NLogN)

이것 또한 시간 복잡도는 NLogN이다. 전체 숫자에 대해 heappush를 통해 힙에 넣고 정렬을 하기 때문에 대상의 갯수인 N에 heap에 넣는 행위에 대한 시간복잡도 LogN이 이루어진다.

 

 

Space Complexity: O(N)

현재 숫자를 카운트하기 위해서 dictionary를 사용했기 때문에 배열 내의 Unique한 숫자의 갯수만큼이 공간복잡도가 될 것이다. 

 

다른 해결 방식

1. 이를 더 간단하게 파이썬의 내장 함수만을 사용한다면?(공식 풀이)

 

from collections import Counter
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]: 
        if k == len(nums):
            return nums
        
        count = Counter(nums)   

        return heapq.nlargest(k, count.keys(), key=count.get)

=> Counter라는 파이썬 내장함수를 통해 배열 내의 숫자를 바로 세고 이를 heapq를 이용하여 가장 큰  k개를 가져오는 방식이다.

 

2. Quickselect(Hoare's selection algorithm)이라는 새로운 알고리즘을 사용해보자!(공식 풀이)

 

우선 간단하게 Quickselect의 idea와 과정에 대해 설명하고 이를 코드에 대입해보자.

 

<주로 사용하는 곳>

Quickselect는 보통 "find kth something" 즉, k개의 무언가를 찾을 때 주로 사용하면 좋은 알고리즘으로 예를 들어 k개의 가장 작은 것, k개의 가장 큰 것, k개의 가장 빈번한 것과 같은 문제에서 주로 사용한다.

 

< Idea >

Idea를 간단하게 소개해본다면 자신이 임의로 pivot을 선택하고 이를 기준으로 왼쪽에는 해당 피봇을 기준으로 기준을 충족시키지 못하는 것들을 두고 오른쪽에는 해당 피봇을 기준으로 기준을 넘는 것을 두는 알고리즘으로 이를 사용하면 해당 pivot의 위치는 모두 정렬했을 때의 위치와 반드시 같을 수밖에 없다. 그러면 해당 pivot이 만약에 n-k위치에 있다고 한다면 해당 pivot을 포함한 뒤에 있는 원소들은 현재 찾는 k개의 원소를 의미하므로 거기서 정렬을 그만두고 리턴하는 것이다.

( 그래서 이번 문제에서 순서 상관없이 원소를 리턴해도 된다고 한 것!!)

 

<과정>

1. Counter를 이용하여 각 원소의 빈도수를 딕셔너리 형태로 저장하고 그 키를 리스트로 저장한다. k를 2라 하자

{1:2, 2:5, 3:6, 4:3, 5:1, 6:4}

unique = [1,2,3,4,5,6]

 

2.임의의 피봇을 선택하고 해당 pivot과 리스트의 맨 마지막에 있는 원소와 바꾸고 이후 left를 stored_index로 지정하자. i는 현재 순회하고 있는 위치로 범위는 left<=i<right이다.

1(left, stored_index,i)

2

3

6

5

4(right)

 

3. left위치에 stored_index로 지정해주고 비교를 시작한다. 

1) 1의 빈도수 < 4의 빈도수

=> 1과 현재 stored_index의 원소와 자리를 바꾸고 stored_index와 i를 증가. ( 1 <-> 1)

1(left)

2(stored_index,i)

3

6

5

4(right)

2) 2의 빈도수 > 4의 빈도수

=> 만약 pivot의 빈도수보다 크다면 그대로 두고 i만을 증가한다.

1(left)

2(stored_index)

3(i)

6

5

4(right)

3) 3의 빈도수 > 4의 빈도수

=> pivot의 빈도수보다 크니까 그대로 두고 i만을 증가한다.

1(left)

2(stored_index)

3

6(i)

5

4(right)

4) 6의 빈도수 > 4의 빈도수

=> pivot의 빈도수보다 크니까 그대로 두고 i만을 증가한다.

1(left)

2(stored_index)

3

6

5(i)

4(right)

5) 5의 빈도수 < 4의 빈도수

=> pivot의 빈도수보다 작으므로 stored_index와 i자리에 있는 원소를 swap하고 stored_index를 증가시킨다.

1(left)

5

3(stored_index)

6

2(i)

4(right)

6) 순회가 끝났으므로 현재 pivot인 4와 stored_index 자리에 있는 원소를 swap한다.

1(left)

5

4(pivot)

6

2

3(right)

{1:2, 5:1, 4:3, 6:4, 2:5, 3:6}

unique = [1,5,4,6,2,3]

=> 이렇게 하면 pivot이었던 4를 기준으로 왼쪽에는 빈도수가 3보다 작은 애들 오른쪽에는 빈도수가 4보다 큰 애들이 나오는 것을 확인할 수 있다.

 

7) n-k와 현재 pivot의 인덱스를 비교하여 n-k가 작으면 right를 pivot 오른쪽으로 크면 left를 pivot왼쪽으로 움직여 앞선 과정을 그대로 반복해주면 된다.  같은 경우 해당 pivot을 포함한 오른쪽 원소 모두가 답이 된다.

(n-k는 뒤에서 k번째로 큰 수의 index를 의미하며 이는 pivot의 index와 n-k와 같아야 함을 알 수 있다.)

 

=> 즉 임의의 원소를 잡아 pivot을 설정하고 이를 기준으로 대소관계를 이용해 분류한 후 해당 pivot의 위치를 이용하여 답을 내는 게 이 알고리즘의 핵심이라고 볼 수 있다.

 

 

이를 코드로 변환하면..

from collections import Counter
class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = Counter(nums)  # counter를 이용하여 각 빈도수를 딕셔너리 형태로 저장
        unique = list(count.keys()) # key만을 unique에 리스트 형태로 저장
        
        def partition(left, right, pivot_index) -> int:
            pivot_frequency = count[unique[pivot_index]] # 해당 pivot의 인덱스를 확인
            unique[pivot_index], unique[right] = unique[right], unique[pivot_index]  
            # pivot을 맨 오른쪽으로 보냄
            
            store_index = left
            for i in range(left, right):
                if count[unique[i]] < pivot_frequency: # 현재의 숫자의 빈도수가 피봇의 빈도수보다 작다면
                    unique[store_index], unique[i] = unique[i], unique[store_index] # store_index와 현재 key를 바꿈
                    store_index += 1   # 증가

            unique[right], unique[store_index] = unique[store_index], unique[right]  # 마지막에 pivot과 저장위치를 바꿈.
            
            return store_index
        
        def quickselect(left, right, k_smallest) -> None:
            if left == right: 
                return
            
            pivot_index = random.randint(left, right)    # 임의의 left와 right사이의 index를 잡음
                              
            pivot_index = partition(left, right, pivot_index) # pivot의 인덱스를 찾아
            
            if k_smallest == pivot_index:  # 만약에 pivot인덱스와 찾으려는 시작점이 같다면 끝
                 return 
            elif k_smallest < pivot_index:  # 만약에 찾으려는 시작점이 피봇보다 더 작다면 왼쪽을 다시 찾음
                quickselect(left, pivot_index - 1, k_smallest)
            else:                         # 만약에 찾으려는 시작점이 피봇 오른쪽에 위치한다면 오른쪽을 quickselect함.
                quickselect(pivot_index + 1, right, k_smallest)
         
        n = len(unique) 
 
        quickselect(0, n - 1, n - k) # left를 시작값, n-1을 끝값, n-k를 내가 끝내고 싶은 지점

        return unique[n - k:] # n-k부터 k개의 원소를 리턴함.

 

문제 링크

https://leetcode.com/problems/top-k-frequent-elements/description/

Top K Frequent Elements - LeetCode