An interesting question I completed awhile ago but I redid it today for fun.

I knew this problem involved heaps in some way, because we have to sort some group of values based on some specific characteristic. In this case, we need to sort based on their distances from the origin, and then somehow get K of the closest elements. This is a big signal for using heaps

With that in mind, it is pretty straight forward. In problems involving heaps, we can solve it with min or max heap. Which one should we choose? We know that with heaps, the top element is always the smallest (or largest in a max-heap).

If we use a min heap, we can store all the elements in it and then remove K times from the heap. This will work because in a min heap, the elements will be sorted from smallest to greatest. Time complexity would be O(n log n) since for each heap operation, it costs O(log n) complexity.

However, if we use a max heap and keep it at a size K, we can reduce this complexity down to O(n log k), where K is the size of the heap. This is because heap operations depend on the size of the heap itself.

Why does max-heap work? It seems a little strange that a max-heap will work for getting the K closest elements because max-heaps store big to small, but we need the K smallest. The trick here lies in keeping a max-heap of size K.

**If we keep it at size K, the largest element in the heap will be at the root of the heap. This means these K elements are actually the K closest. Imagine if we had (1,2,3,4) in our heap and K = 4. If we add 5, this will be the new root element but the heap size > k. So we remove from the heap to maintain the size. When we remove, we remove 5 and are left with (1,2,3,4).**

Similarly, if we add 0, now 4 is no longer part of the closest and it will be removed. Our new heap will have (0,1,2,3), which will be the 4 closest points to the origin.

If we are asked to find the K largest numbers, we could use a min-heap and apply similar concepts. Or if we are asked to find the K’th smallest value in an array, a max-heap will also work

Lastly, the important part is the comparator function. We need to sort based on distance. So all we have to do is calculate their distance and store it as part of the tuple and build a custom comparator. I like to use a lambda function since it makes the code easier and shorter

`class Solution {`

public int[][] kClosest(int[][] points, int k) {

int[][] ans = new int[k][2];

// Sort based on closest

Queue<int[]> heap = new PriorityQueue<>( (a,b) -> {

return b[2] - a[2];

});

// Calculate distance once and store it

for (int[] pair : points){

int dist = (int) (Math.pow(pair[0],2) + Math.pow(pair[1],2));

int[] group = new int[]{pair[0], pair[1], dist};

heap.add(group);

if (heap.size() > k) heap.remove();

}

for (int i = 0; i < ans.length; i++){

int[] group = heap.remove();

int[] pair = new int[]{group[0], group[1]};

ans[i] = pair;

}

return ans;

}

}