Thursday, June 2, 2016

Leetcode: 347. Top K Frequent Elements

Given a non-empty array of integers, return the k most frequent elements.
For example,
Given [1,1,1,2,2,3] and k = 2, return [1,2].
Note: 
  • You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
  • Your algorithm's time complexity must be better than O(n log n), where n is the array's size.
Solution:
Use a min-heap.

Code (Java):
public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        List<Integer> result = new ArrayList<>();
        if (nums == null || nums.length == 0) {
            return result;
        } 
        
        // step 1: count the freq of each word
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            if (map.containsKey(num)) {
                int freq = map.get(num);
                map.put(num, freq + 1);
            } else {
                map.put(num, 1);
            }
        }
        
        // step 2: use a min-heap to get the top k frequencies. 
        PriorityQueue<Pair> pq = new PriorityQueue<>(new MyPQComparator());
        int count = 0;
        for (int word : map.keySet()) {
            int freq = map.get(word);
            Pair pair = new Pair(freq, word);
            if (count < k) {
                pq.offer(pair);
            } else if (pair.freq > pq.peek().freq) {
                pq.poll();
                pq.offer(pair);
            }
            count++;
        }
        
        // step 3: output the result
        for (Pair pair : pq) {
            result.add(pair.word);
        }
        
        return result;
    }
    
    private class MyPQComparator implements Comparator<Pair> {
        @Override
        public int compare(Pair a, Pair b) {
            return a.freq - b.freq;
        }
    }
}

class Pair {
    int word;
    int freq;
    
    Pair(int freq, int word) {
        this.word = word;
        this.freq = freq;
    }
}

No comments:

Post a Comment