Description

Given some points and an origin point in two-dimensional space, find k points which are nearest to the origin.
Return these points sorted by distance, if they are same in distance, sorted by the x-axis, and if they are same in the x-axis, sorted by y-axis.

Example

Example 1:

Input: points = [[4,6],[4,7],[4,4],[2,5],[1,1]], origin = [0, 0], k = 3
Output: [[1,1],[2,5],[4,4]]

Solution

可以对所有的进行排序，但是没必要(只需要keep K个即可，zillow面试就是这么挂的qaq)
注意堆只min heap，需要对所有值取反
时间O(NlogK)

"""
Definition for a point.
class Point:
    def __init__(self, a=0, b=0):
        self.x = a
        self.y = b
"""
import heapq
class Solution:
    """
    @param points: a list of points
    @param origin: a point
    @param k: An integer
    @return: the k closest points
    """
    def dis(self,p1,p2):
        return - ( (p1.x-p2.x)**2 + (p1.y-p2.y)**2 )
    def kClosest(self, points, origin, k):
        # write your code here
        self.heap = []
        heapq.heapify(self.heap)
        for p in points:
            
            if len(self.heap) <k:
                heapq.heappush(self.heap,(self.dis(p,origin),-p.x,-p.y))
            else:
                max_p = heapq.heappop(self.heap)
                if max_p[0] > self.dis(p,origin):
                    push_p = max_p
                # Redundant code(还是直接存x,y heap 已经可以sort)
                # elif max_p[0] == self.dis(p,origin):
                #     if p.x < -max_p[1]:
                #         push_p = (self.dis(p,origin),-p.x,-p.y) 
                #     elif p.x == -max_p[1]:
                #         push_p = (self.dis(p,origin),-p.x,-p.y) if p.y < -max_p[2] else max_p
                #     else:
                #         push_p = max_p
                else:
                    push_p = (self.dis(p,origin),-p.x,-p.y) 
                heapq.heappush(self.heap,push_p)
        res = []
        for i in range(k):
            p=heapq.heappop(self.heap)
            res.append([-p[1],-p[2]])
        res.reverse()
        return res
              

封装版，添加比较函数

"""
Definition for a point.
class Point:
    def __init__(self, a=0, b=0):
        self.x = a
        self.y = b
"""
import heapq
def compare(x,y):
    
    if x.x > y.x:
        return True
    elif x.x == y.x:
        if x.y > y.y:
            return True
        else:
            return False
    else:
        return False

Point.__lt__ = lambda x,y : compare(x,y)

class Solution:
    """
    @param points: a list of points
    @param origin: a point
    @param k: An integer
    @return: the k closest points
    """
    def dis(self,x,y):
        return - ((x.x-y.x)**2+(x.y-y.y)**2)
    def kClosest(self, points, origin, k):
        # write your code here
        self.heap=[]
        heapq.heapify(self.heap)
        for p in points:
            if len(self.heap) <k:
                heapq.heappush(self.heap,(self.dis(p,origin),p))
            else:
                max_p = heapq.heappop(self.heap)
                if max_p[0] > self.dis(p,origin):
                    push_p = max_p
                elif max_p[0] == self.dis(p,origin):
                    if max_p[1] < p:
                        push_p = (self.dis(p,origin),p)
                    else:
                        push_p = max_p
                else:
                    push_p = (self.dis(p,origin),p)
                heapq.heappush(self.heap, push_p)
        res =[]
        for i in range(k):
            dis,p = heapq.heappop(self.heap)
            res.append([p.x,p.y])
        res.reverse()
        return res