最接近原点的 K 个点

题目描述

leetcode-cn.com/problems/k-…

分析

题目找到最接近原点的 K 个点，实际上就找最小 K 个

解题思路

利用大顶堆来找 K 个元素

算法

大顶堆

过程

将所有的点放入堆中，如果堆中元素数量在放入后大于 K，弹出一个

代码

/**
 * @param {number[][]} points
 * @param {number} k
 * @return {number[][]}
 */
var kClosest = function (points, k) {
    const getDistance = (p) => {
        const x = p[0]
        const y = p[1]

        return Math.pow(x, 2) + Math.pow(y, 2)
    }

    const h = new Heap((a, b) => {
        if (!b) return false
        return getDistance(a) < getDistance(b)
    })
    for (const x of points) {
        h.insert(x)

        if (h.size() > k) h.extract()
    }

    return h.heap
};

class Heap {
    constructor(compareFn) {
        this.compareFn = compareFn
        this.heap = []
    }

    getLeftIndex(index) {
        return index * 2 + 1
    }

    getRightIndex(index) {
        return index * 2 + 2
    }

    getParentIndex(index) {
        return Math.floor((index - 1) / 2)
    }

    size() {
        return this.heap.length
    }

    swap(parent, index) {
        const arr = this.heap;
        ;[arr[parent], arr[index]] = [arr[index], arr[parent]]
    }

    isEmpty() {
        return this.size() === 0
    }

    insert(value) {
        const index = this.heap
        this.heap.push(value)
        this.siftUp(index)
    }

    siftUp(index) {
        let parent = this.getParentIndex(index)

        while (index > 0 && this.compareFn(this.heap[parent], this.heap[index])) {
            this.swap(parent, index)
            index = parent
            parent = this.getParentIndex(index)
        }
    }

    extract() {
        if (this.isEmpty()) return
        if (this.size() === 1) return this.heap.pop()

        const removedItem = this.heap[0]
        this.heap[0] = this.heap.pop()
        this.siftDown(0)

        return removedItem
    }

    siftDown(index) {
        let element = index
        const left = this.getLeftIndex(index)
        const right = this.getRightIndex(index)

        if (index < this.size() && this.compareFn(this.heap[element], this.heap[left])) {
            element = left
        }

        if (index < this.size() && this.compareFn(this.heap[element], this.heap[right])) {
            element = right
        }

        if (index !== element) {
            this.swap(element, index)
            this.siftDown(element)
        }
    }

    top() {
        return this.heap[0]
    }
}