K Closest Points to the Origin

Heap

Easy

We have a list ofpoints on the plane. Find theKclosest points to the origin(0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in.)

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.)

Note:

1. 1 <= K <= points.length <= 10000

2. -10000 < points[i][0] < 10000

3. -10000 < points[i][1] < 10000

Solution

Use Max-Heap

Space: O(k) - max heap of size k

Time: O(n logk + k)

import java.util.*;
public class Main {
    public static List<List<Integer>> topK(List<List<Integer>> input, int n, int k){ 
        LinkedList<List<Integer>> ans = new LinkedList<>();
        PriorityQueue<List<Integer>> maxHeap = new PriorityQueue<List<Integer>>(k, new Comparator<List<Integer>>() {
            public int compare(List<Integer> a, List<Integer> b) {
                return b.get(0)*b.get(0) + b.get(1) * b.get(1) - a.get(0) * a.get(0) - a.get(1) * a.get(1);
            }
        });

        for (List<Integer> c: input) {
            maxHeap.offer(c);
            if (!maxHeap.isEmpty() && maxHeap.size() > k) {
                maxHeap.poll();
            }
        }
        while (!maxHeap.isEmpty()) {
            ans.addFirst(new ArrayList<Integer>(maxHeap.poll()));
        }
        return ans;
    }

    public static void main (String[] args){
        List<List<Integer>> input = new ArrayList<>();
        input.add(Arrays.asList(-1,0));
        input.add(Arrays.asList(1,1));
        input.add(Arrays.asList(1,2));
        input.add(Arrays.asList(1,3));
        input.add(Arrays.asList(1,4));
        input.add(Arrays.asList(2,5));
        int n = 6;
        int k = 2;
        System.out.println(topK(input,n,k));
    }

}

Another Max-Heap (with lambda function)

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(
            (p1, p2) -> p2[0] * p2[0] + p2[1] * p2[1] - p1[0] * p1[0] - p1[1] * p1[1]);
        for (int[] p : points) {
            pq.offer(p);
            if (pq.size() > K) {
                pq.poll();
            }
        }
        int size = Math.min(K, pq.size());
        int[][] res = new int[size][2];
        int i = size;
        while (!pq.isEmpty()) {
            res[--i] = pq.poll();
        }
        return res;
    }
}

Use Min-Heap

public static List<List<Integer>> topK(List<List<Integer>> input, int n,int m){
        PriorityQueue<List<Integer>> pq = new PriorityQueue<List<Integer>>(n,
                new Comparator<List<Integer>>(){
                  public int compare (List<Integer> e1,List<Integer> e2){
                      return e1.get(0)*e1.get(0) + e1.get(1)*e1.get(1) - e2.get(0)*e2.get(0) - e2.get(1)*e2.get(1);
                  }
                });
        for(List<Integer> e1:input){
            pq.add(e1);
        }
        List<List<Integer>> result = new ArrayList<>();
        for(int i = 0;i < m && i < n;i++){
            result.add(pq.remove());
        }
        return result;    
    }

Reference

https://www.lintcode.com/problem/k-closest-points/description

From: https://www.youtube.com/watch?v=eaYX0Ee0Kcg

Three solutions to this classical K-th problem. https://leetcode.com/problems/k-closest-points-to-origin/discuss/220235/Java-Three-solutions-to-this-classical-K-th-problem.