sundeepblue
4/26/2014 - 3:48 AM

given a list of points in 2D, and a single reference point, find k nearest neighbors. k closest points. kth

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