given a list of points in 2D, and a single reference point, find k nearest neighbors. k closest points. kth
#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <cstring>
#include <queue>
using namespace std;
struct point {
int x, y;
point(int _x, int _y) : x(_x), y(_y) {}
};
struct mypair {
float dist;
point p;
mypair(float d, point q) : dist(d), p(q) {}
//operator < (mypair& m) { return dist < m.dist; } // not work for priority_queue
};
class Comp {
public:
bool operator () (mypair& p1, mypair& p2) {
return p1.dist < p2.dist;
}
};
float get_distance(point p1, point p2) {
return sqrt( (p1.x - p2.x) * (p1.x - p2.x) + (p1.y - p2.y) * (p1.y - p2.y) );
}
vector<point> KNN(vector<point> &vp, int K, point p) {
vector<point> res;
if(vp.empty()) return res;
if(K >= vp.size()) return vp;
// typedef pair<float, point> pfp;
priority_queue<mypair, vector<mypair>, Comp> dist_queue;
for(int i=0; i<K; i++) {
float d = get_distance(vp[i], p);
dist_queue.push(mypair(d, vp[i]));
}
for(int i=K; i<vp.size(); i++) {
float d = get_distance(vp[i], p);
if(d >= dist_queue.top().dist) continue;
dist_queue.pop();
dist_queue.push(mypair(d, vp[i]));
}
/* // wrong! cannot iterate priority_queue!!!
// http://stackoverflow.com/questions/18492927/can-i-access-elements-of-a-priority-que-using-an-iterator
for(auto p : dist_queue)
res.push_back(p.p);
*/
while(dist_queue.empty() == false) {
res.push_back(dist_queue.top().p);
dist_queue.pop();
}
return res;
}
int main()
{
vector<point> vp = {{1,1}, {2,2}, {1,0}, {0,1}, {3,3},{5,5}};
vector<point> knn = KNN(vp, 5, point(0,0));
for(auto p : knn) cout << p.x << " " << p.y << endl;
}