#include <stdlib.h>
typedef struct {
int dist; // squared distance, same order as the real one
int *point;
} Entry;
static void swap(Entry *a, Entry *b) {
Entry t = *a;
*a = *b;
*b = t;
}
static void siftUp(Entry *h, int i) {
while (i > 0 && h[(i - 1) / 2].dist < h[i].dist) {
swap(&h[(i - 1) / 2], &h[i]);
i = (i - 1) / 2;
}
}
static void siftDown(Entry *h, int n, int i) {
for (;;) {
int max = i, l = 2 * i + 1, r = 2 * i + 2;
if (l < n && h[l].dist > h[max].dist)
max = l;
if (r < n && h[r].dist > h[max].dist)
max = r;
if (max == i)
return;
swap(&h[max], &h[i]);
i = max;
}
}
int **kClosest(int **points, int pointsSize, int *pointsColSize, int k,
int *returnSize, int **returnColumnSizes) {
Entry *pq = malloc((k + 1) * sizeof(Entry)); // max-heap
int n = 0;
for (int i = 0; i < pointsSize; i++) {
int *p = points[i];
pq[n] = (Entry){p[0] * p[0] + p[1] * p[1], p};
siftUp(pq, n++);
if (n > k) {
pq[0] = pq[--n];
siftDown(pq, n, 0);
}
}
int **res = malloc(n * sizeof(int *));
*returnColumnSizes = malloc(n * sizeof(int));
for (int i = 0; i < n; i++) {
res[i] = malloc(2 * sizeof(int));
res[i][0] = pq[i].point[0];
res[i][1] = pq[i].point[1];
(*returnColumnSizes)[i] = 2;
}
*returnSize = n;
free(pq);
return res;
}
#include <cmath>
#include <queue>
#include <utility>
#include <vector>
using namespace std;
class Solution {
public:
vector<vector<int>> kClosest(vector<vector<int>> &points, int k) {
priority_queue<pair<double, vector<int>>> pq;
for (auto &it : points) {
double dist = sqrt(it[0] * it[0] + it[1] * it[1]);
pq.push({dist, it});
if (pq.size() > k)
pq.pop();
}
vector<vector<int>> res;
while (!pq.empty()) {
res.push_back(pq.top().second);
pq.pop();
}
return res;
}
};
package heaps
import "container/heap"
type entry struct {
dist int // squared distance, same order as the real one
point []int
}
type maxHeap []entry
func (h maxHeap) Len() int { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i].dist > h[j].dist }
func (h maxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any) { *h = append(*h, x.(entry)) }
func (h *maxHeap) Pop() any {
old := *h
x := old[len(old)-1]
*h = old[:len(old)-1]
return x
}
func kClosest(points [][]int, k int) [][]int {
pq := &maxHeap{}
for _, p := range points {
heap.Push(pq, entry{p[0]*p[0] + p[1]*p[1], p})
if pq.Len() > k {
heap.Pop(pq)
}
}
res := make([][]int, 0, k)
for _, e := range *pq {
res = append(res, e.point)
}
return res
}
import java.util.PriorityQueue;
class Solution {
private static int dist(int[] p) {
return p[0] * p[0] + p[1] * p[1]; // squared, same order as the real one
}
public int[][] kClosest(int[][] points, int k) {
PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> dist(b) - dist(a));
for (int[] p : points) {
pq.offer(p);
if (pq.size() > k)
pq.poll();
}
return pq.toArray(new int[0][]);
}
}
class MaxHeap {
constructor(key) {
this.key = key;
this.data = [];
}
size() {
return this.data.length;
}
push(val) {
const h = this.data;
h.push(val);
let i = h.length - 1;
while (i > 0) {
const p = (i - 1) >> 1;
if (this.key(h[p]) >= this.key(h[i])) break;
[h[p], h[i]] = [h[i], h[p]];
i = p;
}
}
pop() {
const h = this.data;
const top = h[0];
const last = h.pop();
if (h.length === 0) return top;
h[0] = last;
let i = 0;
for (;;) {
let max = i;
const l = 2 * i + 1,
r = 2 * i + 2;
if (l < h.length && this.key(h[l]) > this.key(h[max])) max = l;
if (r < h.length && this.key(h[r]) > this.key(h[max])) max = r;
if (max === i) return top;
[h[max], h[i]] = [h[i], h[max]];
i = max;
}
}
}
/**
* @param {number[][]} points
* @param {number} k
* @return {number[][]}
*/
var kClosest = function (points, k) {
// squared distance, same order as the real one
const pq = new MaxHeap(([x, y]) => x * x + y * y);
for (const p of points) {
pq.push(p);
if (pq.size() > k) pq.pop();
}
return pq.data;
};
import heapq
from typing import List
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
pq = [] # max-heap via negated squared distance
for x, y in points:
heapq.heappush(pq, (-(x * x + y * y), x, y))
if len(pq) > k:
heapq.heappop(pq)
return [[x, y] for _, x, y in pq]
use std::collections::BinaryHeap;
impl Solution {
pub fn k_closest(points: Vec<Vec<i32>>, k: i32) -> Vec<Vec<i32>> {
// keyed by squared distance, same order as the real one
let mut pq = BinaryHeap::new();
for p in points {
pq.push((p[0] * p[0] + p[1] * p[1], p));
if pq.len() > k as usize {
pq.pop();
}
}
pq.into_iter().map(|(_, p)| p).collect()
}
}