ヒープで流れてくる点数の k 番目に高い値をすぐ答える仕組みを作る

from heapq import heappop
from heapq import heappush


class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.h = []
        for x in nums:
            self.add(x)

    def add(self, val):
        heappush(self.h, val)
        if len(self.h) > self.k:
            heappop(self.h)
        return self.h[0]

class KthLargest {
  constructor(k, nums) {
    this.k = k;
    this.h = [];
    for (const x of nums) {
      this.add(x);
    }
  }

  add(val) {
    const h = this.h;
    h.push(val);
    up(h, h.length - 1);
    if (h.length > this.k) {
      h[0] = h.pop();
      down(h, 0);
    }
    return h[0];
  }
}

function up(h, i) {
  while (i > 0) {
    const p = (i - 1) >> 1;
    if (h[p] <= h[i]) {
      return;
    }
    [h[p], h[i]] = [h[i], h[p]];
    i = p;
  }
}

function down(h, i) {
  const n = h.length;
  while (true) {
    let s = i;
    const l = 2 * i + 1;
    const r = l + 1;
    if (l < n && h[l] < h[s]) {
      s = l;
    }
    if (r < n && h[r] < h[s]) {
      s = r;
    }
    if (s === i) {
      return;
    }
    [h[s], h[i]] = [h[i], h[s]];
    i = s;
  }
}

import "container/heap"

type minHeap []int

func (h minHeap) Len() int {
	return len(h)
}

func (h minHeap) Less(
	i, j int,
) bool {
	return h[i] < h[j]
}

func (h minHeap) Swap(
	i, j int,
) {
	h[i], h[j] = h[j], h[i]
}

func (h *minHeap) Push(x any) {
	*h = append(*h, x.(int))
}

func (h *minHeap) Pop() any {
	old := *h
	x := old[len(old)-1]
	*h = old[:len(old)-1]
	return x
}

type KthLargest struct {
	k int
	h *minHeap
}

func newKthLargest(
	k int, nums []int,
) *KthLargest {
	s := &KthLargest{
		k: k, h: &minHeap{}}
	for _, x := range nums {
		s.Add(x)
	}
	return s
}

func (s *KthLargest) Add(
	val int,
) int {
	heap.Push(s.h, val)
	if s.h.Len() > s.k {
		heap.Pop(s.h)
	}
	return (*s.h)[0]
}