Code › algorithm-study
LeetCode 347 - Top K Frequent Elements
빈도수를 세고 bucket sort 방식으로 상위 k개 원소를 찾는 풀이
빈도수가 높은 원소 k개를 반환하는 Medium 문제다. 정렬로 풀 수도 있지만, 이번 풀이는 빈도수를 index로 쓰는 bucket 배열 방식으로 정리했다.
문제 링크 & 설명
- 문제 링크: 347. Top K Frequent Elements
- 요약: 정수 배열 nums와 정수 k가 주어졌을 때, 가장 자주 등장한 원소 k개를 반환하는 문제.
예를 들어 입력이 이렇게 주어진다고 해보자.
nums = [1, 1, 1, 2, 2, 3]
k = 2
1은 세 번, 2는 두 번 등장하므로 결과는 [1, 2]가 된다. 순서는 문제에서 중요하지 않다.
접근 방법
먼저 hash map으로 각 숫자의 빈도수를 센다.
freq_map = {}
for num in nums:
freq_map[num] = freq_map.get(num, 0) + 1
그 다음 빈도수를 index로 쓰는 buckets 배열을 만든다. nums의 길이가 n이면 어떤 숫자도 n번보다 많이 등장할 수 없으므로, bucket은 n + 1개면 충분하다.
buckets = [[] for i in range(len(nums) + 1)]
for num, freq in freq_map.items():
buckets[freq].append(num)
이제 높은 빈도수부터 내려오면서 원소를 꺼내면 된다. k개를 모은 순간 바로 반환한다.
top_list = []
for i in range(len(nums), -1, -1):
for num in buckets[i]:
top_list.append(num)
if len(top_list) == k:
return top_list
핵심은 값을 정렬하는 대신 빈도수를 위치로 바꾸는 것이다. 빈도수 3인 값은 buckets[3]에 들어가고, 빈도수 2인 값은 buckets[2]에 들어간다. 그래서 뒤에서부터 훑으면 높은 빈도수의 값부터 만난다.
복잡도 분석
-
시간 복잡도: O(n)
- 빈도수 계산에 한 번, bucket 채우기에 한 번, 뒤에서부터 결과를 모으는 데 최대 한 번 순회한다.
-
공간 복잡도: O(n)
- freq_map과 buckets가 입력 크기에 비례하는 공간을 사용한다.
정렬을 사용하면 보통 O(n log n) 비용이 들어가지만, 이 방식은 빈도수의 가능한 범위가 0부터 n까지라는 점을 이용해서 선형 시간으로 처리한다.
최적화 코드
from typing import List
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
freq_map = {}
for num in nums:
freq_map[num] = freq_map.get(num, 0) + 1
buckets = [[] for _ in range(len(nums) + 1)]
for num, freq in freq_map.items():
buckets[freq].append(num)
top_list = []
for i in range(len(nums), -1, -1):
for num in buckets[i]:
top_list.append(num)
if len(top_list) == k:
return top_list
return top_list
요약 및 회고
이 문제는 hash map으로 빈도수를 세는 것까지는 바로 떠올리기 쉽다. 그 다음에 모든 원소를 빈도수 기준으로 정렬하면 풀리지만, 빈도수가 가질 수 있는 최대값이 nums의 길이라는 점을 이용하면 bucket으로 바꿀 수 있고 정렬 없이 시간복잡도 O(n)를 유지할 수 있다.