├── .gitignore ├── .travis.yml ├── LICENSE ├── README.md ├── movingmedian.go └── movingmedian_test.go /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled Object files, Static and Dynamic libs (Shared Objects) 2 | *.o 3 | *.a 4 | *.so 5 | 6 | # Folders 7 | _obj 8 | _test 9 | 10 | # Architecture specific extensions/prefixes 11 | *.[568vq] 12 | [568vq].out 13 | 14 | *.cgo1.go 15 | *.cgo2.c 16 | _cgo_defun.c 17 | _cgo_gotypes.go 18 | _cgo_export.* 19 | 20 | _testmain.go 21 | 22 | *.exe 23 | *.test 24 | *.prof 25 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | 3 | go: 4 | - 1.x 5 | - 1.6 6 | - 1.7.x 7 | - 1.8.x 8 | - master 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2015 JaderDias 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | movingmedian 2 | ========== 3 | 4 | [![GoDoc](https://godoc.org/github.com/JaderDias/movingmedian?status.svg)](https://godoc.org/github.com/JaderDias/movingmedian) 5 | [![Build Status](https://travis-ci.org/JaderDias/movingmedian.svg?branch=master)](https://travis-ci.org/JaderDias/movingmedian) 6 | 7 | Description 8 | ----------- 9 | 10 | Package movingmedian computes the median of a windowed stream of data. 11 | 12 | Installation 13 | ------------ 14 | 15 | This package can be installed with the go get command: 16 | 17 | go get github.com/JaderDias/movingmedian 18 | -------------------------------------------------------------------------------- /movingmedian.go: -------------------------------------------------------------------------------- 1 | // Package movingmedian computes the median of a windowed stream of data. 2 | package movingmedian 3 | 4 | import "container/heap" 5 | 6 | type item struct { 7 | f float64 8 | heapIndex int 9 | } 10 | 11 | type itemHeap []*item 12 | 13 | func (h itemHeap) Len() int { return len(h) } 14 | func (h itemHeap) Swap(i, j int) { 15 | h[i], h[j] = h[j], h[i] 16 | h[i].heapIndex = i 17 | h[j].heapIndex = j 18 | } 19 | 20 | func (h *itemHeap) Push(x interface{}) { 21 | e := x.(*item) 22 | e.heapIndex = len(*h) 23 | *h = append(*h, e) 24 | } 25 | 26 | func (h *itemHeap) Pop() interface{} { 27 | old := *h 28 | n := len(old) 29 | x := old[n-1] 30 | *h = old[0 : n-1] 31 | return x 32 | } 33 | 34 | type minItemHeap struct { 35 | itemHeap 36 | } 37 | 38 | func (h minItemHeap) Less(i, j int) bool { return h.itemHeap[i].f < h.itemHeap[j].f } 39 | 40 | type maxItemHeap struct { 41 | itemHeap 42 | } 43 | 44 | func (h maxItemHeap) Less(i, j int) bool { return h.itemHeap[i].f > h.itemHeap[j].f } 45 | 46 | // MovingMedian computes the moving median of a windowed stream of numbers. 47 | type MovingMedian struct { 48 | queueIndex int 49 | nitems int 50 | queue []item 51 | maxHeap maxItemHeap 52 | minHeap minItemHeap 53 | } 54 | 55 | // NewMovingMedian returns a MovingMedian with the given window size. 56 | func NewMovingMedian(size int) MovingMedian { 57 | m := MovingMedian{ 58 | queue: make([]item, size), 59 | maxHeap: maxItemHeap{}, 60 | minHeap: minItemHeap{}, 61 | } 62 | 63 | heap.Init(&m.maxHeap) 64 | heap.Init(&m.minHeap) 65 | return m 66 | } 67 | 68 | // Push adds an element to the stream, removing old data which has expired from the window. It runs in O(log windowSize). 69 | func (m *MovingMedian) Push(v float64) { 70 | if len(m.queue) == 1 { 71 | m.queue[0].f = v 72 | return 73 | } 74 | 75 | itemPtr := &m.queue[m.queueIndex] 76 | m.queueIndex++ 77 | if m.queueIndex >= len(m.queue) { 78 | m.queueIndex = 0 79 | } 80 | 81 | if m.nitems == len(m.queue) { 82 | minAbove := m.minHeap.itemHeap[0].f 83 | maxBelow := m.maxHeap.itemHeap[0].f 84 | itemPtr.f = v 85 | if itemPtr.heapIndex < m.minHeap.Len() && itemPtr == m.minHeap.itemHeap[itemPtr.heapIndex] { 86 | if v >= maxBelow { 87 | heap.Fix(&m.minHeap, itemPtr.heapIndex) 88 | return 89 | } 90 | 91 | rotate(&m.maxHeap, &m.minHeap, m.maxHeap.itemHeap, m.minHeap.itemHeap, itemPtr) 92 | return 93 | } 94 | 95 | if v <= minAbove { 96 | heap.Fix(&m.maxHeap, itemPtr.heapIndex) 97 | return 98 | } 99 | 100 | rotate(&m.minHeap, &m.maxHeap, m.minHeap.itemHeap, m.maxHeap.itemHeap, itemPtr) 101 | return 102 | } 103 | 104 | m.nitems++ 105 | itemPtr.f = v 106 | if m.minHeap.Len() == 0 || v > m.minHeap.itemHeap[0].f { 107 | heap.Push(&m.minHeap, itemPtr) 108 | rebalance(&m.minHeap, &m.maxHeap) 109 | } else { 110 | heap.Push(&m.maxHeap, itemPtr) 111 | rebalance(&m.maxHeap, &m.minHeap) 112 | } 113 | } 114 | 115 | func rebalance(heapA, heapB heap.Interface) { 116 | if heapA.Len() == (heapB.Len() + 2) { 117 | moveItem := heap.Pop(heapA) 118 | heap.Push(heapB, moveItem) 119 | } 120 | } 121 | 122 | func rotate(heapA, heapB heap.Interface, itemHeapA, itemHeapB itemHeap, itemPtr *item) { 123 | moveItem := itemHeapA[0] 124 | moveItem.heapIndex = itemPtr.heapIndex 125 | itemHeapB[itemPtr.heapIndex] = moveItem 126 | itemHeapA[0] = itemPtr 127 | heap.Fix(heapB, itemPtr.heapIndex) 128 | itemPtr.heapIndex = 0 129 | heap.Fix(heapA, 0) 130 | } 131 | 132 | // Median returns the current value of the median from the window. 133 | func (m *MovingMedian) Median() float64 { 134 | if len(m.queue) == 1 { 135 | return m.queue[0].f 136 | } 137 | 138 | if m.maxHeap.Len() == m.minHeap.Len() { 139 | return (m.maxHeap.itemHeap[0].f + m.minHeap.itemHeap[0].f) / 2 140 | } 141 | 142 | if m.maxHeap.Len() > m.minHeap.Len() { 143 | return m.maxHeap.itemHeap[0].f 144 | } 145 | 146 | return m.minHeap.itemHeap[0].f 147 | } 148 | -------------------------------------------------------------------------------- /movingmedian_test.go: -------------------------------------------------------------------------------- 1 | package movingmedian 2 | 3 | import ( 4 | "math" 5 | "math/rand" 6 | "sort" 7 | "testing" 8 | ) 9 | 10 | func TestUnit(t *testing.T) { 11 | tests := []struct { 12 | name string 13 | windowSize int 14 | data []float64 15 | want []float64 16 | }{ 17 | { 18 | "OneWindowSize", 19 | 1, 20 | []float64{1, 3, 5, 7, 9, 11, math.NaN()}, 21 | []float64{1, 3, 5, 7, 9, 11, math.NaN()}, 22 | }, 23 | { 24 | "OddWindowSize", 25 | 3, 26 | []float64{1, 3, 5, 7, 9, 11}, 27 | []float64{1, 2, 3, 5, 7, 9}, 28 | }, 29 | { 30 | "EvenWindowSize", 31 | 4, 32 | []float64{1, 3, 5, 7, 9, 11}, 33 | []float64{1, 2, 3, 4, 6, 8}, 34 | }, 35 | { 36 | "DecreasingValues", 37 | 4, 38 | []float64{19, 17, 15, 13, 11, 9}, 39 | []float64{19, 18, 17, 16, 14, 12}, 40 | }, 41 | { 42 | "DecreasingIncreasingValues", 43 | 4, 44 | []float64{21, 19, 17, 15, 13, 11, 13, 15, 17, 19}, 45 | []float64{21, 20, 19, 18, 16, 14, 13, 13, 14, 16}, 46 | }, 47 | { 48 | "IncreasingDecreasingValues", 49 | 4, 50 | []float64{11, 13, 15, 17, 19, 21, 19, 17, 15, 13}, 51 | []float64{11, 12, 13, 14, 16, 18, 19, 19, 18, 16}, 52 | }, 53 | { 54 | 55 | "ZigZag", 56 | 4, 57 | []float64{21, 23, 17, 27, 13, 31, 9, 35, 5, 39, 1}, 58 | []float64{21, 22, 21, 22, 20, 22, 20, 22, 20, 22, 20}, 59 | }, 60 | { 61 | 62 | "NewValuesInBetween", 63 | 4, 64 | []float64{21, 21, 19, 19, 21, 21, 19, 19, 19, 19}, 65 | []float64{21, 21, 21, 20, 20, 20, 20, 20, 19, 19}, 66 | }, 67 | { 68 | "SameNumberInBothHeaps3Times", 69 | 4, 70 | []float64{11, 13, 13, 13, 25, 27, 29, 31}, 71 | []float64{11, 12, 13, 13, 13, 19, 26, 28}, 72 | }, 73 | { 74 | "SameNumberInBothHeaps3TimesDecreasing", 75 | 4, 76 | []float64{31, 29, 29, 29, 17, 15, 13, 11}, 77 | []float64{31, 30, 29, 29, 29, 23, 16, 14}, 78 | }, 79 | { 80 | "SameNumberInBothHeaps4Times", 81 | 4, 82 | []float64{11, 13, 13, 13, 13, 25, 27, 29, 31}, 83 | []float64{11, 12, 13, 13, 13, 13, 19, 26, 28}, 84 | }, 85 | } 86 | 87 | for _, test := range tests { 88 | t.Log("test name", test.name) 89 | m := NewMovingMedian(test.windowSize) 90 | for i, v := range test.data { 91 | m.Push(v) 92 | actual := m.Median() 93 | if test.want[i] != actual && !(math.IsNaN(actual) && math.IsNaN(test.want[i])) { 94 | firstElement := 1 + i - test.windowSize 95 | if firstElement < 0 { 96 | firstElement = 0 97 | } 98 | t.Errorf("failed on test %s index %d the median of %f is %f and not %f", 99 | test.name, 100 | i, 101 | test.data[firstElement:1+i], 102 | test.want[i], 103 | actual) 104 | } 105 | } 106 | } 107 | } 108 | 109 | func TestRandom(t *testing.T) { 110 | rangeSize := 1000 111 | for windowSize := 1; windowSize < 50; windowSize++ { 112 | data := getData(rangeSize, windowSize) 113 | intData := make([]int, rangeSize) 114 | for i, v := range data { 115 | intData[i] = int(v) 116 | } 117 | 118 | t.Log("test name random test window size", windowSize) 119 | m := NewMovingMedian(windowSize) 120 | for i, v := range data { 121 | want := median(data, i, windowSize) 122 | 123 | m.Push(v) 124 | actual := m.Median() 125 | if want != actual { 126 | firstElement := 1 + i - windowSize 127 | if firstElement < 0 { 128 | firstElement = 0 129 | } 130 | 131 | t.Errorf("failed on test random window size %d index %d the median of %d is %f and not %f", 132 | windowSize, 133 | i, 134 | intData[firstElement:1+i], 135 | want, 136 | actual) 137 | } 138 | } 139 | } 140 | } 141 | 142 | func Benchmark_10values_windowsize1(b *testing.B) { 143 | benchmark(b, 10, 1) 144 | } 145 | 146 | func Benchmark_100values_windowsize10(b *testing.B) { 147 | benchmark(b, 100, 10) 148 | } 149 | 150 | func Benchmark_10Kvalues_windowsize100(b *testing.B) { 151 | benchmark(b, 10000, 100) 152 | } 153 | 154 | func Benchmark_10Kvalues_windowsize1000(b *testing.B) { 155 | benchmark(b, 10000, 1000) 156 | } 157 | 158 | func benchmark(b *testing.B, numberOfValues, windowSize int) { 159 | data := getData(numberOfValues, windowSize) 160 | 161 | b.ResetTimer() 162 | 163 | for i := 0; i < b.N; i++ { 164 | m := NewMovingMedian(windowSize) 165 | for _, v := range data { 166 | m.Push(v) 167 | m.Median() 168 | } 169 | } 170 | } 171 | 172 | func getData(rangeSize, windowSize int) []float64 { 173 | var data = make([]float64, rangeSize) 174 | var r = rand.New(rand.NewSource(99)) 175 | for i := range data { 176 | data[i] = math.Floor(1000 * r.Float64()) 177 | } 178 | 179 | return data 180 | } 181 | 182 | func median(data []float64, i, windowSize int) float64 { 183 | min := 1 + i - windowSize 184 | if min < 0 { 185 | min = 0 186 | } 187 | 188 | if len(data) == 0 { 189 | return math.NaN() 190 | } 191 | 192 | window := make([]float64, 1+i-min) 193 | copy(window, data[min:i+1]) 194 | 195 | if len(window) == 1 { 196 | return window[0] 197 | } 198 | 199 | sort.Float64s(window) 200 | 201 | k := len(window) / 2 202 | if len(window)%2 == 1 { 203 | return window[k] 204 | } 205 | 206 | return 0.5*window[k-1] + 0.5*window[k] 207 | } 208 | --------------------------------------------------------------------------------