973. K Closest Points to Origin

Medium

Description

給定平面上的點陣列 points，找出距離原點 (0, 0) 最近的 k 個點。距離用歐幾里得距離計算。答案可以任意順序回傳。

Example:

Input: points = [[1,3],[-2,2]], k = 1 Output: [[-2,2]]

Intuition#

核心思路：「最近的 K 個」等價於 Top-K 問題，用大小為 k 的 Max Heap 維護。

距離原點的歐幾里得距離 = sqrt(x^2 + y^2)，比較大小時不需要開根號，直接比 x^2 + y^2
Top-K 最小 → 用大小 k 的 Max Heap，堆頂是目前 k 個中最遠的，新點更近就替換
也可以用排序或 Quick Select 來解

Approaches#

1: Sorting#

概念: 按距離排序，取前 k 個
時間複雜度: O(n log n)
空間複雜度: O(n) (排序所需)

class Solution {
    fun kClosest(points: Array<IntArray>, k: Int): Array<IntArray> {
        points.sortBy { it[0] * it[0] + it[1] * it[1] }
        return points.copyOfRange(0, k)
    }
}

⭐2: Max Heap (大小 k)#

概念: 維護大小為 k 的 Max Heap（按距離排序），遍歷所有點，若距離小於堆頂就替換
時間複雜度: O(n log k)
空間複雜度: O(k)

class Solution {
    fun kClosest(points: Array<IntArray>, k: Int): Array<IntArray> {
        val maxHeap = PriorityQueue<IntArray>(compareByDescending { it[0] * it[0] + it[1] * it[1] })

        for (point in points) {
            maxHeap.offer(point)
            if (maxHeap.size > k) {
                maxHeap.poll()
            }
        }

        return maxHeap.toTypedArray()
    }
}

補充：Quick Select

利用 Quick Select 演算法，平均 O(n) 時間找到第 k 小的距離分界點，將陣列分成前 k 小和其餘。

class Solution {
    fun kClosest(points: Array<IntArray>, k: Int): Array<IntArray> {
        quickSelect(points, 0, points.size - 1, k)
        return points.copyOfRange(0, k)
    }

    private fun quickSelect(points: Array<IntArray>, left: Int, right: Int, k: Int) {
        if (left >= right) return

        val pivotIdx = partition(points, left, right)
        when {
            pivotIdx == k -> return
            pivotIdx < k -> quickSelect(points, pivotIdx + 1, right, k)
            else -> quickSelect(points, left, pivotIdx - 1, k)
        }
    }

    private fun partition(points: Array<IntArray>, left: Int, right: Int): Int {
        val pivot = dist(points[right])
        var i = left
        for (j in left until right) {
            if (dist(points[j]) <= pivot) {
                points[i] = points[j].also { points[j] = points[i] }
                i++
            }
        }
        points[i] = points[right].also { points[right] = points[i] }
        return i
    }

    private fun dist(point: IntArray): Int = point[0] * point[0] + point[1] * point[1]
}

時間複雜度: 平均 O(n)，最壞 O(n^2)
空間複雜度: O(1)（原地操作）

🔑 Takeaways#

Pattern: Top-K 最小用大小 k 的 Max Heap；Top-K 最大用大小 k 的 Min Heap
關鍵技巧: 比較歐幾里得距離時不需要開根號，直接比平方和即可；Quick Select 可做到平均 O(n)