Amazon OA2 K-Nearest Point C++

题目描述：

给定N个坐标Point，每个Point实例有x-坐标和y-坐标。题目要求函数返回离原点最近的k个坐标。

思路：

这道题和找第k大或第k小的题目的思路基本相同，就是在遍历所有Point的同时，维护一个size为k的max—heap，一旦发现size为k+1，我们就把max-heap头上最大的元素移出heap，因为这里的heap是max-heap，所以heap头部的元素比heap里其他的元素都要比heap里的其他元素离原点远。这样使得heap里的元素是到目前为止里原点最近的k的点。

复杂度分析

时间复杂度：O(NlogK)
因为需要遍历所有元素，每次遍历一个元素的同时，还要在耗费logk的时间来维护heap。

空间复杂度： O(K)
heap的size 是k

// Example program
#include <iostream>
#include <string>
#include <algorithm>
#include <queue>
#include <math.h>  
#include <vector>
using namespace std;

struct Point { 
    double x;
    double y; 
    Point(double a, double b) {
        x = a;
        y = b;
    }
};

double getDistance(Point a, Point b) {
    return sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
}
typedef bool (*comp)(Point, Point);
Point global_origin = Point(0,0);
bool compare(Point a, Point b)
{
   return (getDistance(a, global_origin)< getDistance(b, global_origin));
}

vector<Point> Solution(vector<Point> &array, Point origin, int k) {
    global_origin = Point(origin.x, origin.y);
    priority_queue<Point, std::vector<Point>, comp> pq(compare);
    vector<Point> ret;
    for (int i = 0; i < array.size(); i++) {
        Point p = array[i];
        pq.push(p);
        if (pq.size() > k)
            pq.pop();
    }
    int index = 0;
    while (!pq.empty()){
        Point p = pq.top();
        ret.push_back(p);
        pq.pop();
    }
    return ret;
}



int main()
{
   Point p1 = Point(4.5, 6.0);
   Point p2 = Point(4.0, 7.0);
   Point p3 = Point(4.0, 4.0);
   Point p4 = Point(2.0, 5.0);
   Point p5 = Point(1.0, 1.0);
   vector<Point> array = {p1, p2, p3, p4, p5};
   int k = 2;
   Point origin = Point(0.0, 0.0);
   vector<Point> ans = Solution(array, origin, k);
   for (int i = 0; i < ans.size(); i++) {
       cout << i << ": " << ans[i].x << "," << ans[i].y << endl;
   }
   //cout << getDistance(p1, p2) << endl;
}