Median Tracker With Two Heaps

package main

import (
	"container/heap"
	"fmt"
)

type maxHeap []int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(x any)        { *h = append(*h, x.(int)) }
func (h *maxHeap) Pop() any {
	old := *h
	n := len(old)
	value := old[n-1]
	*h = old[:n-1]
	return value
}

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
	n := len(old)
	value := old[n-1]
	*h = old[:n-1]
	return value
}

type MedianTracker struct {
	lower maxHeap
	upper minHeap
}

func (t *MedianTracker) Add(value int) {
	if t.lower.Len() == 0 || value <= t.lower[0] {
		heap.Push(&t.lower, value)
	} else {
		heap.Push(&t.upper, value)
	}
	t.rebalance()
}

func (t *MedianTracker) rebalance() {
	for t.lower.Len() > t.upper.Len()+1 {
		heap.Push(&t.upper, heap.Pop(&t.lower))
	}
	for t.upper.Len() > t.lower.Len() {
		heap.Push(&t.lower, heap.Pop(&t.upper))
	}
}

func (t *MedianTracker) Median() float64 {
	if t.lower.Len() == 0 {
		return 0
	}
	if t.lower.Len() == t.upper.Len() {
		return float64(t.lower[0]+t.upper[0]) / 2
	}
	return float64(t.lower[0])
}

func main() {
	tracker := &MedianTracker{}
	for _, value := range []int{7, 3, 9, 1, 5} {
		tracker.Add(value)
		fmt.Printf("after %d: median=%.1f lower=%d upper=%d\n", value, tracker.Median(), tracker.lower.Len(), tracker.upper.Len())
	}
}