leetcode-go/leetcode/medium/215.kth-largest-element-in-an-array.go

func findKthLargest(nums []int, k int) int {
	pq := NewFixedPQ(k)
	for i := range nums {
		pq.Enqueue(nums[i])
	}
	return pq.Peek()
}

type FixedPQ struct { // Small root heap XD
	Size int
	PQ   []int
}

func NewFixedPQ(size int) FixedPQ {
	return FixedPQ{size, make([]int, 1)}
}

func (pq *FixedPQ) Enqueue(x int) {
	if pq.Len() < pq.Size {
		pq.PQ = append(pq.PQ, x)
		pq.swim(pq.Len())
	} else if pq.Peek() < x {
		pq.PQ[1] = x
		pq.sink(1)
	}
}

func (pq *FixedPQ) Dequeue() int {
	peek := pq.PQ[1]
	pq.Swap(1, pq.Len())
	pq.PQ = pq.PQ[:pq.Len()]
	pq.sink(1)
	return peek
}

func (pq FixedPQ) Len() int           { return len(pq.PQ) - 1 }
func (pq FixedPQ) Less(i, j int) bool { return pq.PQ[i] < pq.PQ[j] }
func (pq FixedPQ) Swap(i, j int)      { pq.PQ[i], pq.PQ[j] = pq.PQ[j], pq.PQ[i] }

func (pq FixedPQ) Peek() int {
	return pq.PQ[1]
}

func (pq FixedPQ) swim(i int) {
	for 1 < i && pq.Less(i, i/2) {
		pq.Swap(i, i/2)
		i /= 2
	}
}

func (pq FixedPQ) sink(i int) {
	for 2*i <= pq.Size {
		j := 2*i
		if (j < pq.Size && pq.Less(j+1, j)) {
			j++ // Compare with the smaller child
		}
		if !pq.Less(j, i) {
			break // If ith <= jth, stop sinking
		}
		pq.Swap(i, j)
		i = j
	}
}