83 lines
1.7 KiB
Go
83 lines
1.7 KiB
Go
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
|
|
}
|
|
}
|