[Leetcode 347] Top K Frequent Elements - 해시테이블, 우선순위큐

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

LeetCode - The World's Leading Online Programming Learning Platform

 

Problems

빈도순으로 k개의 엘리먼트를 추출하라. O(nlogn)보다 더 빠른 시간복잡도로 구현해라. 

 

입력: nums = [1, 1, 1, 2, 2, 3, 4], k = 2

출력: [1, 2]

 

설명:

1이 총 3개, 2가 총 2개, 3과 4는 각각 1개의 빈도수를 가지고 있다. 빈도순대로 정렬하면 1,2 순서대로 큰 빈도값을 갖고 있다.

 

---

전략 1. 빈도 수를 해시맵에 저장하고 빈도 순으로 엘리먼트 추출

예를들어 nums가 [1,1,1,2,2,3,4] 라고 하자. 먼저 각 엘리먼트의 빈도를 추출한다.

{ 1 : 3 , 2 : 2 , 3 : 1 , 4 : 1} 로 맵에 저장시킨다. 우리가 원하는 것은 빈도수대로 k 만큼 키를 뽑아내야 한다. 빈도수가 3인 것은 1 element.

2인것은 2 element. 1인것은 3,4 element 이다. map에는 {빈도수,element} 형태로 만들고 빈도수대로 element 를 뽑아내면 될 것이다.

코드를 보면 다음과 같이 이해가 될 것이다.

 

Map<Integer, Integer> frequentMap = new HashMap<>();
for (int num : nums) {
    frequentMap.put(num, frequentMap.getOrDefault(num, 0) + 1);
}

Map<Integer, List<Integer>> buckets = new HashMap<>();
for (Integer key : frequentMap.keySet()) {
    int frequency = frequentMap.get(key);
    List<Integer> elems = buckets.getOrDefault(frequency, new ArrayList<>());
    elems.add(key);
    buckets.put(frequency, elems);
}

 

다음 이 buckets라는 해시맵을 k번만큼 추출하면 되므로 result를 k크기만큼 선언하고 idx를 선언해서 k의 빈도만큼 저장했는지 체킹하는 용도로 사용한다.

int[] result = new int[k];
int idx = 0;

 

빈도수는 최대 nums의 길이만큼이다. 그러므로 거기서부터 하나씩 줄여가며 빈도수가 해시맵에 존재하면 추출하여 result에 넣는 방식으로 진행한다.

 

for (int frequency = nums.length; frequency >= 0; frequency--) {
    if (buckets.containsKey(frequency)) {
        List<Integer> elems = buckets.get(frequency);
        for (Integer elem : elems) {
            if (idx == k) {
                break;
            }
            result[idx] = elem;
            idx++;
        }
    }
    if (idx == k) {
        break;
    }
}
return result;

 

전체 코드는 다음과 같다. 총 시간복잡도는 O(2n) 으로 O(nlogn) 보단 사소한 차이로 빠르긴 하다.

 

public int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> frequentMap = new HashMap<>();
    for (int num : nums) {
        frequentMap.put(num, frequentMap.getOrDefault(num, 0) + 1);
    }

    Map<Integer, List<Integer>> buckets = new HashMap<>();
    for (Integer key : frequentMap.keySet()) {
        int frequency = frequentMap.get(key);
        List<Integer> elems = buckets.getOrDefault(frequency, new ArrayList<>());
        elems.add(key);
        buckets.put(frequency, elems);
    }

    int[] result = new int[k];
    int idx = 0;

    for (int frequency = nums.length; frequency >= 0; frequency--) {
        if (buckets.containsKey(frequency)) {
            List<Integer> elems = buckets.get(frequency);
            for (Integer elem : elems) {
                if (idx == k) {
                    break;
                }
                result[idx] = elem;
                idx++;
            }
        }
        if (idx == k) {
            break;
        }
    }
    return result;

}

 

---

전략 2. 우선순위 큐를 이용한 풀이

 

위 풀이랑 차이점은 우선순위큐에 {요소,빈도수}를 넣는다음 빈도수 기준으로 우선순위큐에 삽입(내림차순으로) 한다. 그다음에 k번 우선순위큐에서 추출하면 훨씬 간단하게 풀 수 있을 것이다.

public int[] topKFrequent(int[] nums, int k) {
    Map<Integer,Integer> frequencyMap = new HashMap<>();
    for(int n : nums){
        frequencyMap.put(n,frequencyMap.getOrDefault(n,0)+1);
    }

    PriorityQueue<int[]> pq = new PriorityQueue<>((a,b)->b[1]-a[1]);
    for(int elem : frequencyMap.keySet()){
        pq.add(new int[]{elem,frequencyMap.get(elem)});
    }

    int[] result = new int[k];
    for(int i = 0; i <k; i++){
        result[i] = pq.poll()[0];
    }
    return result;
}