NeetCodeのSolutionを書いていく

#

Heap / Priority Queue

#

問題

#

2次元配列の points と整数 kが与えられます。 X-Y平面上の点を構成する座標を表します。 points[i] = [xi, yi] です。

[0, 0]に近いk個の点を返しましょう。 距離はユークリッド距離とします。 答えの順番は任意です。

例

#

Input: points = [[0,2],[2,2]], k = 1 Output: [[0,2]]

Input: points = [[0,2],[2,0],[2,2]], k = 2 Output: [[0,2],[2,0]]

制約

#

• 1 <= k <= points.length <= 1000 -100 <= points[i][0], points[i][1] <= 100
• 計算量：O(nlogk)
• メモリ：O(k)

メモ

#

heaqpを使うのがポイント。

import heapq
from typing import List


class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        minHeap = []

        # 計算量：O(n)
        for x, y in points:
            dist = (x**2) + (y**2)
            minHeap.append([dist, x, y])

        # 計算量：O(n)
        heapq.heapify(minHeap)
        res = []

        # 計算量O(k*logn)
        while k > 0:
            dist, x, y = heapq.heappop(minHeap)
            res.append([x, y])
            k -= 1

        return res