[파이썬 알고리즘]973.leetcode - 원점에 k번째로 가까운 점

https://leetcode.com/problems/k-closest-points-to-origin/

K Closest Points to Origin - LeetCode

 

문제) 평면상에 points 목록이 있을 때, 원점 (0, 0)에서 k번 가까운 점 목록을 순서대로 출력하라. 평면상 두 점의 거리는 유클리드 거리로 한다.

 

 

풀이)

import heapq
import math


def cal_dist(point):
    return math.sqrt(point[0] * point[0] + point[1] * point[1])


def k_closest_builtin(points, k):
    dists = []
    heap = []
    for point in points:
        dist = cal_dist(point)
        heapq.heappush(heap, dist)
        dists.append(dist)

    kth_dist = [heapq.heappop(heap) for _ in range(k)][-1]
    return [points[idx] for idx, dist in enumerate(dists) if dist <= kth_dist]

print(k_closest_builtin([[3,3],[5,1],[-2,4]], 2))