Rolling median in go

main.go

package main

import (
	"fmt"
	"math"
	"sort"
	"time"
)

// Data holds Value & Time when value was observed
type Data struct {
	Value float64
	Time  time.Time
}

// SortedFloatSlice assumes elements are sorted
type SortedFloatSlice []float64

// Insert into slice maintaing the sort order
func (f SortedFloatSlice) Insert(value float64) SortedFloatSlice {
	i := sort.SearchFloat64s(f, value)
	n := append(f, 0)
	copy(n[i+1:], n[i:])
	n[i] = value
	return n
}

// Delete from slice maintaing the sort order
func (f SortedFloatSlice) Delete(value float64) SortedFloatSlice {
	i := sort.SearchFloat64s(f, value)
	n := append(f[:i], f[i+1:]...)
	return n
}

// Median of the slice
func (f SortedFloatSlice) Median() float64 {
	if len(f)%2 == 1 {
		return f[len(f)/2]
	}
	return (f[len(f)/2] + f[len(f)/2-1]) / 2
}

// FloatQueue is FIFO implementation
type FloatQueue []float64

// Push to queue
func (fs FloatQueue) Push(item float64) FloatQueue {
	return append(fs, item)
}

// Pop from queue
func (fs FloatQueue) Pop() (item float64, queue FloatQueue) {
	return fs[0], fs[1:]
}

// FixedSizeSortedQueue remembers the insert order & keeps last N elements sorted
type FixedSizeSortedQueue struct {
	Queue  FloatQueue
	Sorted SortedFloatSlice
}

// Push will insert a new element, remove oldest & keep the slice sorted
func (s *FixedSizeSortedQueue) Push(item float64) {
	removed, newQueue := s.Queue.Pop()
	newQueue = newQueue.Push(item)
	newSorted := s.Sorted.Delete(removed)
	newSorted = newSorted.Insert(item)
	s.Queue = newQueue
	s.Sorted = newSorted
}

// NewFixedSizeSortedQueue initialises the queue of fixed order
func NewFixedSizeSortedQueue(values []float64) FixedSizeSortedQueue {
	sorted := make([]float64, len(values))
	copy(sorted, values)
	sort.Float64s(sorted)
	return FixedSizeSortedQueue{
		Queue:  values,
		Sorted: sorted,
	}
}

func movingMedian(series []Data, period int) []Data {
	res := make([]Data, len(series))

	initialWindow := make([]float64, 1, period)
	initialWindow[0] = 0 // we need one extra element that will be ignored
	for _, v := range series[:(period - 1)] {
		initialWindow = append(initialWindow, v.Value)
	}
	queue := NewFixedSizeSortedQueue(initialWindow)

	median := float64(0)
	for i, v := range series {
		if i < period-1 {
			median = math.NaN()
		} else {
			queue.Push(v.Value)
			median = queue.Sorted.Median()
		}

		res[i] = Data{
			Value: median,
			Time:  v.Time,
		}
	}

	return res
}

func mk(in string) time.Time {
	s, err := time.Parse(time.RFC3339, in)
	if err != nil {
		panic(err)
	}
	return s
}

func main() {
	series := []Data{
		{10, mk("2019-01-01T00:00:00Z")},
		{20, mk("2019-01-01T00:00:00Z")},
		{30, mk("2019-01-01T00:00:00Z")},
		{50, mk("2019-01-01T00:00:00Z")},
	}
	fmt.Printf("%v\n", movingMedian(series, 2))
}