package kthlargestelementinstream

import (
	"container/heap"
)

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

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

func (h IntHeap) Peek() int {
	return h[0]
}

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

type KthLargest struct {
	k       int
	largest IntHeap
}

func Constructor(k int, nums []int) KthLargest {
	result := KthLargest{k, IntHeap{}}
	heap.Init(&result.largest)

	for _, n := range nums {
		result.Add(n)
	}

	return result
}

func (this *KthLargest) Add(val int) int {
	heap.Push(&this.largest, val)
	if this.largest.Len() > this.k {
		heap.Pop(&this.largest)
	}

	return this.largest.Peek()
}