package main

import (
    "container/heap"
    "fmt"
)

// MinHeap structure (heap of larger half)
type MinHeap []int

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] } // Min-heap: smaller values at top
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}
func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// MaxHeap structure (heap of smaller half, but as max-heap)
type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] } // Max-heap: larger values at top
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}
func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// MedianFinder struct using two heaps
type MedianFinder struct {
    maxHeap *MaxHeap // For the smaller half
    minHeap *MinHeap // For the larger half
}

// NewMedianFinder initializes the MedianFinder
func NewMedianFinder() MedianFinder {
    maxHeap := &MaxHeap{}
    minHeap := &MinHeap{}
    heap.Init(maxHeap)
    heap.Init(minHeap)
    return MedianFinder{
        maxHeap: maxHeap,
        minHeap: minHeap,
    }
}

// insertNum inserts a number into the stream
func (mf *MedianFinder) insertNum(num int) {
    // Insert into maxHeap (smaller half) first
    if mf.maxHeap.Len() == 0 || num <= (*mf.maxHeap)[0] {
        heap.Push(mf.maxHeap, num)
    } else {
        heap.Push(mf.minHeap, num)
    }

    // Balance the heaps so that their sizes differ by at most 1
    if mf.maxHeap.Len() > mf.minHeap.Len()+1 {
        heap.Push(mf.minHeap, heap.Pop(mf.maxHeap))
    } else if mf.minHeap.Len() > mf.maxHeap.Len() {
        heap.Push(mf.maxHeap, heap.Pop(mf.minHeap))
    }
}

// findMedian returns the median of the current stream
func (mf *MedianFinder) findMedian() float64 {
    if mf.maxHeap.Len() == mf.minHeap.Len() {
        // Even number of elements, return the average of the two middle elements
        return float64((*mf.maxHeap)[0]+(*mf.minHeap)[0]) / 2.0
    }
    // Odd number of elements, return the root of the max-heap (larger half)
    return float64((*mf.maxHeap)[0])
}