Heap / Priority Queue

class MinHeap<T> {
  private items: T[] = [];
  constructor(private compare: (a: T, b: T) => number) {}
  get size(): number { return this.items.length; }
  push(item: T): void {
    this.items.push(item);
    let i = this.items.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(this.items[parent], this.items[i]) <= 0) break;
      [this.items[parent], this.items[i]] = [this.items[i], this.items[parent]];
      i = parent;
    }
  }
  pop(): T | undefined {
    const top = this.items[0];
    const last = this.items.pop();
    if (this.items.length && last !== undefined) {
      this.items[0] = last;
      let i = 0;
      for (;;) {
        const left = 2 * i + 1, right = left + 1;
        let best = i;
        if (left < this.items.length && this.compare(this.items[left], this.items[best]) < 0) best = left;
        if (right < this.items.length && this.compare(this.items[right], this.items[best]) < 0) best = right;
        if (best === i) break;
        [this.items[i], this.items[best]] = [this.items[best], this.items[i]];
        i = best;
      }
    }
    return top;
  }
}

function kClosest(points: number[][], k: number): number[][] {
  const dist = (p: number[]): number => p[0] * p[0] + p[1] * p[1];
  const heap = new MinHeap<number[]>((a, b) => dist(a) - dist(b));
  for (const p of points) heap.push(p);
  const res: number[][] = [];
  for (let i = 0; i < k && heap.size > 0; i++) res.push(heap.pop()!);
  return res;
}

class MinHeap
  def initialize(&compare)
    @items = []
    @compare = compare
  end

  def size = @items.size

  def push(item)
    @items << item
    i = @items.size - 1
    while i.positive?
      parent = (i - 1) / 2
      break if @compare.call(@items[parent], @items[i]) <= 0
      @items[parent], @items[i] = @items[i], @items[parent]
      i = parent
    end
    self
  end

  def pop
    return nil if @items.empty?
    top = @items[0]
    last = @items.pop
    return top if @items.empty?

    @items[0] = last
    i = 0
    loop do
      left = (2 * i) + 1
      right = left + 1
      best = i
      best = left if left < @items.size && @compare.call(@items[left], @items[best]).negative?
      best = right if right < @items.size && @compare.call(@items[right], @items[best]).negative?
      break if best == i
      @items[i], @items[best] = @items[best], @items[i]
      i = best
    end
    top
  end
end

def k_closest(points, k)
  dist = ->(p) { (p[0] * p[0]) + (p[1] * p[1]) }
  heap = MinHeap.new { |a, b| dist.call(a) - dist.call(b) }
  points.each { |p| heap.push(p) }
  Array.new([k, points.size].min) { heap.pop }
end