LC295 - 数据流的中位数

package Heap

import (
    "cmp"
    "container/heap"
)

// EasyMaxHeap is a generic type.
// The elements should be basic data types.
type EasyMaxHeap[T cmp.Ordered] []T

func (h EasyMaxHeap[T]) Len() int {
    return len(h)
}

func (h EasyMaxHeap[T]) Less(i, j int) bool {
    return h[i] > h[j]
}

func (h EasyMaxHeap[T]) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

func (h *EasyMaxHeap[T]) Push(x interface{}) {
    *h = append(*h, x.(T))
}

func (h *EasyMaxHeap[T]) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func (h *EasyMaxHeap[T]) Top() interface{} {
    if h.Len() == 0 {
        return nil
    } else {
        return (*h)[0]
    }
}

type MedianFinder struct {
    maxHeap *EasyMaxHeap[int]
    minHeap *EasyMinHeap[int]
}

func Constructor() MedianFinder {
    medianFinder := &MedianFinder{
        maxHeap: new(EasyMaxHeap[int]),
        minHeap: new(EasyMinHeap[int]),
    }
    heap.Init(medianFinder.maxHeap)
    heap.Init(medianFinder.minHeap)
    return *medianFinder
}

func (this *MedianFinder) AddNum(num int) {
    if this.maxHeap.Len() == 0 || num <= this.maxHeap.Top().(int) {
        heap.Push(this.maxHeap, num)
        if this.maxHeap.Len()-this.minHeap.Len() > 1 {
            heap.Push(this.minHeap, heap.Pop(this.maxHeap))
        }
    } else {
        heap.Push(this.minHeap, num)
        if this.maxHeap.Len() < this.minHeap.Len() {
            heap.Push(this.maxHeap, heap.Pop(this.minHeap))
        }
    }
}

func (this *MedianFinder) FindMedian() float64 {
    if this.maxHeap.Len() == this.minHeap.Len() {
        return float64(this.maxHeap.Top().(int)+this.minHeap.Top().(int)) / 2
    } else {
        return float64(this.maxHeap.Top().(int))
    }
}

核心思想就是大根堆保存较小的一半，小根堆保存较大的一半。