import heapq

#Solution 1: array sorting, Time worst O(n^2) average O(nlogn), Space O(n)
def kClosest(points, k):
    points.sort(key = lambda p: p[0]*p[0] + p[1]*p[1])
    return points[:k]

#Solution 2, quick select, Time worst O(n^2) average O(n), Space worst O(n) average O(logn)
def kClosestQuick(points, k) :
    n =  len(points)
    low = 0
    high = n - 1
    while (low <= high) :
        mid = partition(points, low, high)
        if mid == k:
            break
        if (mid < k): 
            low = mid + 1 
        else :
            high = mid - 1
    return points[:k]

#Partition, Time worst O(n^2) average O(n), Space O(1)
def partition(points, low, high) :
    pivot = points[low]
    while (low < high) :
        while low < high and compare(points[high], pivot) >= 0: 
            high -= 1
        points[low] = points[high]
        while low < high and compare(points[low], pivot) <= 0: 
            low += 1
        points[high] = points[low]
    points[low] = pivot
    return low

#Solution 3: Priorty queue, Time O(nlogn), Space O(n), n is number of points 
def kClosestPQ(points, k):
    pq = []
    for p in points:
        dist = p[0]*p[0] + p[1]*p[1]
        heapq.heappush(pq, (-dist, p))
        if len(pq) > k:
            heapq.heappop(pq)     
    return [point[1] for point in pq]

#Compare based on Euclidean distance, Time O(1), Space O(1)
def compare(a, b) :
    return (a[0]*a[0] + a[1]*a[1]) - (b[0]*b[0] + b[1]*b[1])


points = [[1,3], [3,2], [4, 5], [2, 4], [2, 3], [3, 7]]
k = 2

#Solution 1, Array sorting
res = kClosest(points, k)
print("sorting: ")
print(res)

#Solution 2, quick select
res1 = kClosestQuick(points,k)
print("quick select: ")
print(res1)

#Solution 3, Priority Queue
res2 = kClosestPQ(points, k)
print("priority queue: ")
print(res2)