Problem

Implement a data structure, provide two interfaces:

1. add(number). Add a new number in the data structure.
2. 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]

思路

public class Solution {

    PriorityQueue<Integer> heap;
    int cap;

    public Solution(int k) {
        heap = new PriorityQueue<>();
        cap = k;
    }

    public void add(int num) {
        heap.add(num);
        if (heap.size() > cap) {
            heap.poll();
        }
    }

    public List<Integer> topk() {
        List<Integer> res = new ArrayList<>();
        for (Integer num: heap) {
            res.add(num);
        }
        Collections.sort(res, Collections.reverseOrder());
        return res;
    }
}

python

import heapq

class Solution:

    # @param {int} k an integer
    def __init__(self, k):
        self.k = k
        self.nums = []


    # @param {int} num an integer
    def add(self, num):
        heapq.heappush(self.nums, num)
        if len(self.nums) > self.k:
            heapq.heappop(self.nums)


    # @return {int[]} the top k largest numbers
    def topk(self):
        return heapq.nlargest(self.k, self.nums)