973. K Closest Points to Origin - Medium

前往題目

這題之前寫過心得，照搬一下

心得

• 這題非常妙，用到的資料結構也不一樣，只要知道如何使用就很簡單
• 使用MinHeap或是MaxHeap就可以完美解決
• 使用MaxHeap的時間複雜度更好一點
• PriorityQueue會根據Comparator來排序，排序時間是logN(因為使用heap，樹結構)，N是項數

思路 (以MinHeap為例子)

1. 新建一個MinHeap
2. 把point都放入minHeap
3. 然後就取出k個回傳，因為MinHeap是由小到大排序，使用Heap的原因是可以在插入數值時可以logN插入，使用Array會是N次

以下使用PriorityQueue來模擬MinHeap

Code

MinHeap做法(NlogN)

public int[][] kClosest(int[][] points, int k) {
    // Use min heap
    PriorityQueue<int[]> q = new PriorityQueue<>((a, b) ->
        Integer.compare(
            (a[0] * a[0] + a[1] * a[1]),
            (b[0] * b[0] + b[1] * b[1])
        )
    );

    // Compare and add to form a min heap
    for(int[] point : points){
        q.add(point);
    }

    // Storage for the result arrays
    int[][] res = new int[k][2];

    // Collect result until k elements
    for(int i = 0; i < k; ++i){
        int[] tmp = q.poll(); // Pop the head, which is the closest point to the origin
        res[i][0] = tmp[0];
        res[i][1] = tmp[1];
    }

    return res;
}

MaxHeap作法(NlogK)

只改了compare的順序，以及q超過k size就去掉head，因為head永遠是最大的，但我們不需要

public int[][] kClosest(int[][] points, int k) {
    PriorityQueue<int[]> q = new PriorityQueue<>((a, b) ->
        Integer.compare(
            (b[0] * b[0] + b[1] * b[1]),
            (a[0] * a[0] + a[1] * a[1])
        )
    ); //only this is changed (swapped)
    for (int[] point : points) {
        q.add(point);
        //remove when size increase k
        if (q.size() > k) {
            q.remove();
        }
    }
    int[][] ans = new int[k][2];
    for (int i = 0; i < k; i++) {
        int[] cur = q.poll();
        ans[i][0] = cur[0];
        ans[i][1] = cur[1];
    }
    return ans;
}