鲸落的博客

Brain-dumping.

Top K Algorithms

背景

去年面试Facebook时,被问到了一道Top K的问题,这类的问题解法往往具有共通性。今天在LeetCode上又遇到了一道类似的问题,于是网上搜了一下解题思路,摘录留下备用。

JDK里有一个强大的集合类 – PriorityQueue,掌握了它的用法便能迎刃而解这类问题。

值得注意的是PriorityQueue没有继承自List,不具有随机访问第i个元素的方法。

1
E get(int index);

Given a non-empty array of integers, return the k most frequent elements.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class Pair {
    int num;
    int count;

    Pair(int num, int count) {
        this.num=num;
        this.count=count;
    }
}

public class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        //count the frequency for each element
        HashMap<Integer, Integer> map = new HashMap<Integer, Integer>();
        for(int num: nums){
            if (map.containsKey(num)) {
                map.put(num, map.get(num)+1);
            } else {
                map.put(num, 1);
            }
        }

        // create a min heap
        PriorityQueue<Pair> queue = new PriorityQueue<>(new Comparator<Pair>() {
            public int compare(Pair a, Pair b) {
                return a.count - b.count;
            }
        });

        //maintain a heap of size k.
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            Pair p = new Pair(entry.getKey(), entry.getValue());
            queue.offer(p);
            if(queue.size() > k) {
                queue.poll();
            }
        }

        //get all elements from the heap
        List<Integer> result = new ArrayList<Integer>();
        while(queue.size() > 0) {
            result.add(queue.poll().num);
        }
        //reverse the order
        Collections.reverse(result);

        return result;
    }
}

Comments