Compute the K Closest Stars

EPI 11.4

Description:

Given n stars, find k stars which are closest to the Earth.

Example:

Note

Idea:

1. if RAM is enough, find kth largest using quick select.
2. use max heap.

Code below use max heap. (not compiled or ran.)

Code:

struct Star{
    bool operator<(const Star& rhs) const{
        return distance() < rhs.distance();
    }
    double distance() const{
        return sqrt(x*x + y*y+ z*z);
    }
    double x, y, z;
};

vector<Star> findKClosestStars(int k, vector<Star> stars){
    priority_queue<Star> closestHeap;

    for(int i=0; i<k; i++){
        closestHeap.push(stars[i]);
    }

    for(int i=k; i<stars.size(); i++){
        if(stars[i]<closestHeap.top()){
            closestHeap.push(stars[i]);
            closestHeap.pop();
        }
    }

    vector<Star> result;
    while(closestHeap.empty()){
        result.push_back(closestHeap.top());
        closestHeap.pop();
    }
    return result;
}