NeetCode 150 [Heap / Priority Queue]medium:K Closest Points to Origin

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