973. K Closest Points to Origin

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        """
        main idea: keep a max heap size of k, push each point into it with distance to origin and its coordinate  
        tc O(NlgK) sc O(K)
        """
        pq = []
        for x,y in points:
            heapq.heappush(pq,(-x*x-y*y,x,y))
            if len(pq)> k:
                heapq.heappop(pq)
        return [[x,y] for _,x,y in pq]

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        """tc O(N) sc O(1)
        """
        self.sort(points,k,0,len(points)-1)
        return points[:k]
    
    def sort(self,A,k,l,r):
        idx = self.partition(A,l,r)
        if idx == k-1:
            return 
        elif idx < k :
            self.sort(A,k,idx+1,r)
        else:
            self.sort(A,k,l,idx-1)
        
        
    def partition(self,A,l,r)->int:
        pivot_idx = random.randint(l,r)
        A[r],A[pivot_idx] =  A[pivot_idx], A[r]
        pivot = A[r]
        for i in range(l,r):
            if A[i][0]*A[i][0] + A[i][1]*A[i][1] < pivot[0]*pivot[0]+pivot[1]*pivot[1]:
                A[i], A[l] = A[l], A[i]
                l += 1 
        A[l], A[r] = A[r], A[l]
        return l