K Closest Points to Origin

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        arr = []
        for i, p in enumerate(points):
            heappush(arr, (-p[0]**2-p[1]**2, i))
            if i >= k:
                heappop(arr)
        return [points[x[1]] for x in arr]