blob: 20ebc8e95a42cab767486be5d9eef5d05ff8aec5 [file] [log] [blame]
// Copyright 2021 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package stats
import (
"math"
)
// IntervalTree is a structure used to make a calculation of running medians quick.
// The structure is described in the Section 8/Appendix of the paper.
type IntervalTree struct {
d int
vals []int
}
// NewIntervalTree creates a new IntervalTree of depth d.
func NewIntervalTree(d int) *IntervalTree {
if d < 0 {
panic("invalid depth")
}
return &IntervalTree{
d: d,
vals: make([]int, (1<<(d+1))-1),
}
}
// walk is a generic function on interval trees to add or remove elements.
func (it *IntervalTree) walk(v float64, update int) {
v = math.Abs(v)
mid, inc := 0.5, 0.25
idx := 0
// Update the levels in the tree.
for i := 0; i <= it.d; i++ {
it.vals[idx] += update
idx = idx*2 + 1
if v > mid {
mid += inc
idx += 1
} else {
mid -= inc
}
inc /= 2.
}
}
// Insert puts an element in an interval tree.
func (it *IntervalTree) Insert(v float64) {
it.walk(v, 1)
}
// Remove removes an element from an interview tree.
func (it *IntervalTree) Remove(v float64) {
it.walk(v, -1)
}
// Median returns the current median.
func (it *IntervalTree) Median() float64 {
// If empty, special case and return 0.
numElements := it.NumElements()
if numElements == 0 {
return 0
}
l, u := 0., 1.
k := int(math.Ceil(float64(numElements) / 2.))
for i := 0; i < len(it.vals); {
j := 2*i + 1
if j >= len(it.vals) {
break
}
if it.vals[i] == k {
kf := float64(k)
a, b := float64(it.vals[j])/kf, float64(it.vals[j+1])/kf
x := (l + (l+u)/2.) / 2.
y := (u + (l+u)/2.) / 2.
return (a*x + b*y) / (a + b)
}
if v := it.vals[j]; v >= k {
i = j
u = (l + u) / 2.
} else {
k -= v
i = j + 1
l = (l + u) / 2.
}
}
return (u-l)/2. + l
}
// NumElements returns the number of elements in the tree.
func (it *IntervalTree) NumElements() int {
return it.vals[0]
}