make removal O(log n)

Previously we were iterating over all the elements in the heap looking
for one to remove, making updates O(windowSize).  This
tracks the index of elements so that removal becomes O(log windowSize)
instead.

benchmark                              old ns/op     new ns/op     delta
Benchmark_10values_windowsize1         4090          1262          -69.14%
Benchmark_100values_windowsize10       46783         23361         -50.07%
Benchmark_10Kvalues_windowsize100      5384446       3378225       -37.26%
Benchmark_10Kvalues_windowsize1000     7133066       3972929       -44.30%

movingmedian.go

@@ -5,20 +5,24 @@ import (
 	"math"
 )
 
-type float64HeapInterface interface {
-	heap.Interface
-	Data() []float64
+type elt struct {
+	f   float64
+	idx int
 }
 
-type float64Heap []float64
+type float64Heap []*elt
 
 func (h float64Heap) Len() int { return len(h) }
 func (h float64Heap) Swap(i, j int) {
 	h[i], h[j] = h[j], h[i]
+	h[i].idx = i
+	h[j].idx = j
 }
 
 func (h *float64Heap) Push(x interface{}) {
-	*h = append(*h, x.(float64))
+	e := x.(*elt)
+	e.idx = len(*h)
+	*h = append(*h, e)
 }
 
 func (h *float64Heap) Pop() interface{} {
@@ -29,44 +33,29 @@ func (h *float64Heap) Pop() interface{} {
 	return x
 }
 
-func (h float64Heap) Data() []float64 {
-	return []float64(h)
-}
-
-func removeFromHeap(h float64HeapInterface, x float64) bool {
-	for i, v := range h.Data() {
-		if v == x {
-			heap.Remove(h, i)
-			return true
-		}
-	}
-
-	return false
-}
-
 type minFloat64Heap struct {
 	float64Heap
 }
 
-func (h minFloat64Heap) Less(i, j int) bool { return h.float64Heap[i] < h.float64Heap[j] }
+func (h minFloat64Heap) Less(i, j int) bool { return h.float64Heap[i].f < h.float64Heap[j].f }
 
 type maxFloat64Heap struct {
 	float64Heap
 }
 
-func (h maxFloat64Heap) Less(i, j int) bool { return h.float64Heap[i] > h.float64Heap[j] }
+func (h maxFloat64Heap) Less(i, j int) bool { return h.float64Heap[i].f > h.float64Heap[j].f }
 
 type MovingMedian struct {
-	size    int
-	queue   []float64
+	idx     int
+	nelts   int
+	queue   []elt
 	maxHeap maxFloat64Heap
 	minHeap minFloat64Heap
 }
 
 func NewMovingMedian(size int) MovingMedian {
 	m := MovingMedian{
-		size:    size,
-		queue:   make([]float64, 0),
+		queue:   make([]elt, size),
 		maxHeap: maxFloat64Heap{},
 		minHeap: minFloat64Heap{},
 	}
@@ -77,46 +66,60 @@ func NewMovingMedian(size int) MovingMedian {
 }
 
 func (m *MovingMedian) Push(v float64) {
-	m.queue = append(m.queue, v)
 
-	if m.minHeap.Len() == 0 ||
-		v > m.minHeap.float64Heap[0] {
-		heap.Push(&m.minHeap, v)
-	} else {
-		heap.Push(&m.maxHeap, v)
-	}
+	if m.nelts >= len(m.queue) {
+		old := &m.queue[m.idx]
 
-	if len(m.queue) > m.size {
-		outItem := m.queue[0]
-		m.queue = m.queue[1:len(m.queue)]
-		if outItem >= m.minHeap.float64Heap[0] {
-			removeFromHeap(&m.minHeap, outItem)
+		if old.f >= m.minHeap.float64Heap[0].f {
+			heap.Remove(&m.minHeap, old.idx)
 		} else {
-			removeFromHeap(&m.maxHeap, outItem)
+			heap.Remove(&m.maxHeap, old.idx)
 		}
 	}
 
+	m.queue[m.idx] = elt{f: v}
+	e := &m.queue[m.idx]
+
+	m.nelts++
+	m.idx++
+
+	if m.idx >= len(m.queue) {
+		m.idx = 0
+	}
+
+	if m.minHeap.Len() == 0 ||
+		v > m.minHeap.float64Heap[0].f {
+		heap.Push(&m.minHeap, e)
+	} else {
+		heap.Push(&m.maxHeap, e)
+	}
+
 	if m.maxHeap.Len() > (m.minHeap.Len() + 1) {
-		moveItem := heap.Pop(&m.maxHeap).(float64)
+		moveItem := heap.Pop(&m.maxHeap)
 		heap.Push(&m.minHeap, moveItem)
 	} else if m.minHeap.Len() > (m.maxHeap.Len() + 1) {
-		moveItem := heap.Pop(&m.minHeap).(float64)
+		moveItem := heap.Pop(&m.minHeap)
 		heap.Push(&m.maxHeap, moveItem)
 	}
 }
 
-func (m MovingMedian) Median() float64 {
+func (m *MovingMedian) Median() float64 {
 	if len(m.queue) == 0 {
 		return math.NaN()
 	}
 
-	if (len(m.queue) % 2) == 0 {
-		return (m.maxHeap.float64Heap[0] + m.minHeap.float64Heap[0]) / 2
+	wsize := m.nelts
+	if m.nelts > len(m.queue) {
+		wsize = len(m.queue)
+	}
+
+	if (wsize % 2) == 0 {
+		return (m.maxHeap.float64Heap[0].f + m.minHeap.float64Heap[0].f) / 2
 	}
 
 	if m.maxHeap.Len() > m.minHeap.Len() {
-		return m.maxHeap.float64Heap[0]
+		return m.maxHeap.float64Heap[0].f
 	}
 
-	return m.minHeap.float64Heap[0]
+	return m.minHeap.float64Heap[0].f
 }