leetcode 973. K Closest Points to Origin

目录

一、问题分析

二、代码实现

1、排序

2、大顶堆

3、快速选择

---

 

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

给定平面上的点的列表，返回前k个距离原点最近的点的集合。

 

一、问题分析

测试用例：

Example 1:
Input: points = [[1,3],[-2,2]], K = 1
Output: [[-2,2]]
Explanation: 
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest K = 1 points from the origin, so the answer is just [[-2,2]].

Example 2:
Input: points = [[3,3],[5,-1],[-2,4]], K = 2
Output: [[3,3],[-2,4]]
(The answer [[-2,4],[3,3]] would also be accepted.)

返回的集合不考虑点的顺序；1<=k<=points.length;

二、代码实现

1、排序

时间为O(nlogn)

class Solution {
    
    //排序
    public int[][] kClosest1(int[][] points, int K) {
        int N = points.length;
        int[] dists = new int[N];
        for (int i=0; i<N; ++i) {
            dists[i] = getDistance(points[i]);
        }

        Arrays.sort(dists);
        int distK = dists[K-1];

        int[][] result = new int[K][2];
        int t = 0;
        for (int i=0; i<N; ++i) {
            if (getDistance(points[i]) <= distK) {
                result[t++] = points[i];
            }
        }
        
        return result;
    }
    private int getDistance(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }
    
}

另一种写法

class Solution {

    //排序
    public int[][] kClosest2(int[][] points, int K) {
        Arrays.sort(points, new PointCompare());
        int[][] res = new int[K][2];
        for(int i = 0; i < K; i++) res[i] = points[i];
        return res;
    }
    class PointCompare implements Comparator<int[]>{
       public int compare(int[] a, int[]b){
            int distA = a[0]*a[0] + a[1]*a[1];
            int distB = b[0]*b[0] + b[1]*b[1];

            return distA < distB ? -1 : 1;
        }
    }

}

2、大顶堆

时间为O(nlogk)

import java.util.Collection;
class Solution {
    public int[][] kClosest3(int[][] points, int K) {    //points[n][2]
        PriorityQueue<Integer> heap = new PriorityQueue<>((n1, n2) -> n2 - n1); //大顶堆，元素为点到原点的距离的平方
        HashMap<Integer, Integer> map = new HashMap<>();    //点的下标 <-> 点到原点的距离的平方
        for (int i=0; i<points.length; i++) {
            int distanceSquare = points[i][0]*points[i][0] + points[i][1]*points[i][1];
            heap.add(distanceSquare);
            map.put(i, distanceSquare);
            
            if (heap.size() >= K+1) {
                int value = heap.poll();
                Collection<Integer> valueSet = map.values();
                valueSet.remove(value);
            }
        }
        
        int[][] result = new int[K][2];
        //Iterator<Integer, Integer> itr = map.iterator();
        //error: wrong number of type arguments: required 1
        //Iterator<Map.Entry<Integer, Integer>> itr = map.iterator();
        //error: cannot find symbol
        Iterator<Map.Entry<Integer, Integer>> itr = map.entrySet().iterator();
        for (int i=0; i<K; i++) {
            int key = itr.next().getKey();
            result[i][0] = points[key][0];
            result[i][1] = points[key][1];
        }
        
        return result;
    }
}

3、快速选择

平均情况的时间为O(n)，最坏情况的时间为O(n^2)

class Solution {
    
    public int[][] kClosest4(int[][] points, int K) {
        int N = points.length;
        int[] dists = new int[N];
        for (int i=0; i<N; ++i) {
            dists[i] = getDistance(points[i]);
        }
        
        int distK = findKth(dists, K);
        
        int[][] result = new int[K][2];
        int t = 0;
        for (int i=0; i<N; ++i) {
            if (getDistance(points[i]) <= distK) {
                result[t++] = points[i];
            }
        }
        
        return result;
    }
    private int findKth(int[] dists, int k) {
        return findKthLargest(dists, dists.length-k+1);
    }
    public int findKthLargest(int[] nums, int k) {
        shuffle(nums);
        return findK(nums, nums.length - k, 0, nums.length - 1);
    }
    private int findK(int[] nums, int k, int l, int r) {
        int pivotIndex = partition(nums, l, r);
 
        if (k < pivotIndex)
            return findK(nums, k, l, pivotIndex - 1);
        else if (k > pivotIndex)
            return findK(nums, k, pivotIndex + 1, r);
        else // k == pivotIndex
            return nums[k];
    }
    private int partition(int[] a, int lo, int hi) {
        int pivotIdx = lo; // bad practice
        int pivotVal = a[pivotIdx];
 
        // put pivot last
        swap(a, pivotIdx, hi);
        int storeIdx = lo;
 
        for (int i = lo; i < hi; i++) {
            if (a[i] <= pivotVal) {
                swap(a, i, storeIdx);
                storeIdx++;
            }
        }
        swap(a, storeIdx, hi);  //storeIdx+1是第一个大于参照元素的元素，storeIdx为参照元素
        return storeIdx;
    }
    private void shuffle(int[] a) {
        for (int i = a.length - 1; i >= 0; i--) {
            int j = (int) (Math.random() * (i + 1)); // j=0..i
            swap(a, i, j);
        }
    }
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

}

 

参考：

https://leetcode.com/problems/k-closest-points-to-origin/discuss/220235/Java-Three-solutions-to-this-classical-K-th-problem.（这篇文章总结得很好）