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)를 유지할 수 있다.