[LintCode] Top k Largest Numbers II

Problem

Implement a data structure, provide two interfaces:

add(number). Add a new number in the data structure.
topk(). Return the top k largest numbers in this data structure. k is given when we create the data structure.

Example

s = new Solution(3);
>> create a new data structure.
s.add(3)
s.add(10)
s.topk()
>> return [10, 3]
s.add(1000)
s.add(-99)
s.topk()
>> return [1000, 10, 3]
s.add(4)
s.topk()
>> return [1000, 10, 4]
s.add(100)
s.topk()
>> return [1000, 100, 10]

Tags

Heap Priority Queue

Solution

public class Solution {
    /*
    * @param k: An integer
    */
    Queue<Integer> pq;
    int size;
    public Solution(int k) {
        // do intialization if necessary
        pq = new PriorityQueue<Integer>();
        size = k;
    }

    /*
     * @param num: Number to be added
     * @return: nothing
     */
    public void add(int num) {
        // write your code here
        if (pq.size() < size) pq.offer(num);
        else if (pq.peek() < num) {
            pq.poll();
            pq.offer(num);
        }
        return;
    }

    /*
     * @return: Top k element
     */
    public List<Integer> topk() {
        // write your code here
        Iterator it = pq.iterator();
        List<Integer> res = new ArrayList<>();
        while (it.hasNext()) {
            res.add((Integer) it.next());
        }
        Collections.sort(res, Collections.reverseOrder());
        return res;
    }
}