[leetcode] Top K Frequent Elements

문제

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

Top K Frequent Elements - LeetCode

내 풀이

class Solution {
    public class Element {
        
        public int number;
        public int counter;
        
        Element(int number, int counter) {
            this.number = number;
            this.counter = counter;
        }
        
        
    }
    public int[] topKFrequent(int[] nums, int k) {
        int[] answer = new int[k];
        PriorityQueue<Element> pq = new PriorityQueue<>((x, y) -> Integer.compare(y.counter, x.counter));
        Map<Integer, Integer> counter = new HashMap<>();
        
        for(int num : nums) {
            counter.put(num, counter.getOrDefault(num, 0) + 1);
        }
        
        counter.entrySet().stream().forEach(e-> pq.add(new Element(e.getKey(), e.getValue())));
        
        
        for(int i=0; i<k; i++) {
            answer[i] = pq.poll().number;
        }
        
        return answer;
    }
}

알게 된 것

- PriorityQueue 는 Queue의 구현클래스다. 그래서 아래처럼 하는게 낫겠다.

Queue<Element> pq = new PriorityQueue<>();

- PQ의 역순 (최대 힙) 은 객체 생성 시 인자에 lambda 함수 (x1, x2) -> Integer.Compare(x2,x1) 을 넣어주면된다.

 

피드백

        // init heap 'the less frequent element first'
        Queue<Integer> heap = new PriorityQueue<>(
            (n1, n2) -> count.get(n1) - count.get(n2));
		
        // 2. keep k top frequent elements in the heap
        // O(N log k) < O(N log N) time
        for (int n: count.keySet()) {
          heap.add(n);
          if (heap.size() > k) heap.poll();    
        }

        // 3. build an output array
        // O(k log k) time
        int[] top = new int[k];
        for(int i = k - 1; i >= 0; --i) {
            top[i] = heap.poll();
        }

답지를 보니 나의 방식보다 조금 달랐다.

나는 전부 PQ에 Max Heap 형태로 넣은 다음, K개만큼 poll 했는데

 

답지는 PQ의 size가 k보다 커질 때마다 1개씩 가장 작은 카운터를 지닌 수를 빼버렸다.

공간복잡도의 측면에서 Heap의 Size는 K보다 커지지 않도록 할 수 있으므로 위 방법이 더 효율적인듯 싶다.

그리고 PQ의 comparator 람다함수도 count.get 을 참조함으로써 나처럼 클래스생성을 할 필요도 없게 하였다.