Amazon OA2 K-Nearest Point C++

Hetian發表於2019-05-09
RESTC++

題目描述：

給定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;
}

相關文章