package findmedianfromdatastream

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 MedianFinder struct {
	// Positive min-heap of top half of values.
	top *IntHeap
	// Negative min-heap of bottom half of values.
	bottom *IntHeap
}

func Constructor() MedianFinder {
	result := MedianFinder{
		top:    &IntHeap{},
		bottom: &IntHeap{},
	}

	heap.Init(result.top)
	heap.Init(result.bottom)
	return result
}

func (m *MedianFinder) AddNum(num int) {
	if m.bottom.Len() > m.top.Len() {
		if -m.bottom.Peek() <= num {
			heap.Push(m.top, num)
		} else {
			val := heap.Pop(m.bottom).(int)
			heap.Push(m.top, -val)
			heap.Push(m.bottom, -num)
		}
	} else if m.bottom.Len() < m.top.Len() {
		if m.top.Peek() >= num {
			heap.Push(m.bottom, -num)
		} else {
			val := heap.Pop(m.top).(int)
			heap.Push(m.bottom, -val)
			heap.Push(m.top, num)
		}
	} else {
		if m.top.Len() == 0 {
			heap.Push(m.top, num)
		} else {
			if median := m.FindMedian(); float64(num) > median {
				heap.Push(m.top, num)
			} else {
				heap.Push(m.bottom, -num)
			}
		}
	}
}

func (m *MedianFinder) FindMedian() float64 {
	if m.bottom.Len() > m.top.Len() {
		return float64(-m.bottom.Peek())
	} else if m.bottom.Len() < m.top.Len() {
		return float64(m.top.Peek())
	} else {
		return float64(m.top.Peek()-m.bottom.Peek()) / 2.0
	}
}