K Closest Points to Origin

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        auto dist = [](const vector<int>& p) {
            return (long long)p[0] * p[0] + (long long)p[1] * p[1];
        };
        priority_queue<pair<long long, int>> pq; // (dist, idx)
        for (int i = 0; i < (int)points.size(); ++i) {
            pq.push({dist(points[i]), i});
            if ((int)pq.size() > k) pq.pop();
        }
        vector<vector<int>> ans;
        while (!pq.empty()) { ans.push_back(points[pq.top().second]); pq.pop(); }
        return ans;
    }
};

import "container/heap"

type pointHeap struct {
    points [][]int
    dists  []int64
}

func (h pointHeap) Len() int           { return len(h.dists) }
func (h pointHeap) Less(i, j int) bool { return h.dists[i] > h.dists[j] }
func (h pointHeap) Swap(i, j int) {
    h.dists[i], h.dists[j] = h.dists[j], h.dists[i]
    h.points[i], h.points[j] = h.points[j], h.points[i]
}
func (h *pointHeap) Push(x interface{}) {
    e := x.([2]interface{})
    h.dists = append(h.dists, e[0].(int64))
    h.points = append(h.points, e[1].([]int))
}
func (h *pointHeap) Pop() interface{} {
    n := len(h.dists)
    d, p := h.dists[n-1], h.points[n-1]
    h.dists = h.dists[:n-1]
    h.points = h.points[:n-1]
    return [2]interface{}{d, p}
}

func kClosest(points [][]int, k int) [][]int {
    h := &pointHeap{}
    for _, p := range points {
        d := int64(p[0])*int64(p[0]) + int64(p[1])*int64(p[1])
        heap.Push(h, [2]interface{}{d, p})
        if h.Len() > k {
            heap.Pop(h)
        }
    }
    ans := make([][]int, 0, k)
    for h.Len() > 0 {
        e := heap.Pop(h).([2]interface{})
        ans = append(ans, e[1].([]int))
    }
    return ans
}

from typing import List
import heapq

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap: List = []
        for x, y in points:
            d = x * x + y * y
            heapq.heappush(heap, (-d, [x, y]))
            if len(heap) > k:
                heapq.heappop(heap)
        return [p for _, p in heap]

class MaxHeap {
    constructor() { this.h = []; }
    size() { return this.h.length; }
    top() { return this.h[0]; }
    push(v) {
        this.h.push(v);
        let i = this.h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[p][0] < this.h[i][0]) { [this.h[p], this.h[i]] = [this.h[i], this.h[p]]; i = p; }
            else break;
        }
    }
    pop() {
        const top = this.h[0];
        const last = this.h.pop();
        if (this.h.length) {
            this.h[0] = last;
            let i = 0, n = this.h.length;
            while (true) {
                const l = 2 * i + 1, r = 2 * i + 2;
                let best = i;
                if (l < n && this.h[l][0] > this.h[best][0]) best = l;
                if (r < n && this.h[r][0] > this.h[best][0]) best = r;
                if (best === i) break;
                [this.h[best], this.h[i]] = [this.h[i], this.h[best]];
                i = best;
            }
        }
        return top;
    }
}

function kClosest(points, k) {
    const pq = new MaxHeap();
    for (const p of points) {
        const d = p[0] * p[0] + p[1] * p[1];
        pq.push([d, p]);
        if (pq.size() > k) pq.pop();
    }
    const ans = [];
    while (pq.size()) ans.push(pq.pop()[1]);
    return ans;
}

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(
            (a, b) -> Long.compare((long) b[0] * b[0] + (long) b[1] * b[1],
                                   (long) a[0] * a[0] + (long) a[1] * a[1])
        );
        for (int[] p : points) {
            pq.offer(p);
            if (pq.size() > k) pq.poll();
        }
        int[][] ans = new int[k][];
        for (int i = 0; i < k; i++) ans[i] = pq.poll();
        return ans;
    }
}

class MaxHeap {
    private h: [number, number[]][] = [];
    size(): number { return this.h.length; }
    push(v: [number, number[]]): void {
        this.h.push(v);
        let i = this.h.length - 1;
        while (i > 0) {
            const p = (i - 1) >> 1;
            if (this.h[p][0] < this.h[i][0]) { [this.h[p], this.h[i]] = [this.h[i], this.h[p]]; i = p; }
            else break;
        }
    }
    pop(): [number, number[]] {
        const top = this.h[0];
        const last = this.h.pop()!;
        if (this.h.length) {
            this.h[0] = last;
            let i = 0;
            const n = this.h.length;
            while (true) {
                const l = 2 * i + 1, r = 2 * i + 2;
                let best = i;
                if (l < n && this.h[l][0] > this.h[best][0]) best = l;
                if (r < n && this.h[r][0] > this.h[best][0]) best = r;
                if (best === i) break;
                [this.h[best], this.h[i]] = [this.h[i], this.h[best]];
                i = best;
            }
        }
        return top;
    }
}

function kClosest(points: number[][], k: number): number[][] {
    const pq = new MaxHeap();
    for (const p of points) {
        const d = p[0] * p[0] + p[1] * p[1];
        pq.push([d, p]);
        if (pq.size() > k) pq.pop();
    }
    const ans: number[][] = [];
    while (pq.size()) ans.push(pq.pop()[1]);
    return ans;
}