导航

求两个数组中两个数和前K小

Posted on 2013-10-09 10:54  wpzhongyq  阅读(157)  评论(0)    收藏  举报

给定两个数组,求a1[i]+a2[j]的前k小个值。

vector<int> FirstKMax(const vector<int> &a1, const vector<int> &a2, int k){
    priority_queue<int> maxHeap;
    for(size_t i=0; i<a1.size(); i++){
        for(size_t j=0; j<a2.size(); j++){
            if(maxHeap.size()<k)
                maxHeap.push(a1[i]+a2[j]);
            else{
                if(maxHeap.top()>a1[i]+a2[j]){
                    maxHeap.pop();
                    maxHeap.push(a1[i]+a2[j]);
                }
                else
                    break;
            }
        }
    }
    vector<int> res;
    while(!maxHeap.empty()){
        res.push_back(maxHeap.top());
        maxHeap.pop();
    }
    reverse(res.begin(),res.end());
    return res;
}
View Code