Top K nearest points to the origin

There are n points in the xy-coordinate system. Find the top K nearest points to the origin.

{(1,2), (-1,2), (3.-3),...}

The idea for this problem is to use the max_heap(yes, the max heap, NOT the min heap) with size K, and define the compare function.

// top K nearest points to the origin
// a problem from mitbbs.com

struct Point {
    double x;
	double y;
	Point(double xx, double yy): x(xx), y(yy){};
};

double distance(const Point &p) {
    return p.x*p.x + p.y*p.y;
}

struct max_heap_compare {
	bool operator()(const Point &a, const Point &b) const {
		return distance(a) < distance(b);
	}
};

class Solution {
public:
	std::vector topKNearest(std::vector &pts, int k) {
	    std::vector res;
		if (pts.empty()) return res;
		std::priority_queue, max_heap_compare> max_heap;
		auto p = pts.begin();
		while(int(max_heap.size()) < k && p!= pts.end()) {
			Point a_pt(p->x, p->y);
			max_heap.push(a_pt);
			++p;
		}
		while(p!=pts.end()) {
			auto& top_pt = max_heap.top();
			if(distance(top_pt) > distance(*p))
			{
				Point ap(p->x, p->y);
				max_heap.pop();
				max_heap.push(ap);
			}
			++p;			
		}
		
		while(!max_heap.empty()) {
			Point p = max_heap.top();
			max_heap.pop();
			res.push_back(p);
		}
		return res;
	}
};


你可能感兴趣的:(Algorithm)