feat: Add running median. (#3)

* feat: Add running median.

* Update docs.

* Implement solution based on heaps.

* Improve readability of code.

Author: goloroden

README.md

@@ -7,6 +7,7 @@ codingcircle contains the examples from the YouTube Coding Cirle.
 So far, the following exercises have been covered:
 
 - [Remove k-th last element](./removekthlastelement/) – removing the k-th last element from a single-linked list
+- [Running median](./runningmedian/) – calculating the running median of a sequence of numbers
 - [Trie](./trie/) – implementing an autocomplete feature using a trie
 
 ## Running quality assurance

runningmedian/documentation.go

@@ -0,0 +1,3 @@
+// Package runningmedian provides an algorithm to calculate the running median
+// across a sequence of numbers.
+package runningmedian

runningmedian/max_heap.go

@@ -0,0 +1,33 @@
+package runningmedian
+
+import "container/heap"
+
+type MaxHeap []int
+
+func NewMaxHeap() *MaxHeap {
+	h := &MaxHeap{}
+	heap.Init(h)
+	return h
+}
+
+func (h MaxHeap) Len() int {
+	return len(h)
+}
+func (h MaxHeap) Less(i, j int) bool {
+	return h[i] > h[j]
+}
+func (h MaxHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+}
+
+func (h *MaxHeap) Push(value any) {
+	*h = append(*h, value.(int))
+}
+
+func (h *MaxHeap) Pop() any {
+	oldHeap := *h
+	valueIndex := len(oldHeap) - 1
+	value := oldHeap[valueIndex]
+	*h = oldHeap[:valueIndex]
+	return value
+}

runningmedian/median.go

@@ -0,0 +1,56 @@
+package runningmedian
+
+import (
+	"container/heap"
+)
+
+func getMedian(minHeap *MinHeap, maxHeap *MaxHeap) float64 {
+	if minHeap.Len() > maxHeap.Len() {
+		return float64((*minHeap)[0])
+	}
+	if minHeap.Len() < maxHeap.Len() {
+		return float64((*maxHeap)[0])
+	}
+	return (float64((*minHeap)[0]) + float64((*maxHeap)[0])) / 2.0
+}
+
+func add(value int, minHeap *MinHeap, maxHeap *MaxHeap) {
+	if minHeap.Len()+maxHeap.Len() < 1 {
+		heap.Push(minHeap, value)
+		return
+	}
+
+	median := getMedian(minHeap, maxHeap)
+	if float64(value) > median {
+		heap.Push(minHeap, value)
+	} else {
+		heap.Push(maxHeap, value)
+	}
+}
+
+func rebalance(minHeap *MinHeap, maxHeap *MaxHeap) {
+	if minHeap.Len() > maxHeap.Len()+1 {
+		root := heap.Pop(minHeap).(int)
+		heap.Push(maxHeap, root)
+	} else if maxHeap.Len() > minHeap.Len()+1 {
+		root := heap.Pop(maxHeap).(int)
+		heap.Push(minHeap, root)
+	}
+}
+
+func GetMedians(values []int) []float64 {
+	minHeap := NewMinHeap()
+	maxHeap := NewMaxHeap()
+
+	medians := []float64{}
+
+	for _, value := range values {
+		add(value, minHeap, maxHeap)
+		rebalance(minHeap, maxHeap)
+
+		median := getMedian(minHeap, maxHeap)
+		medians = append(medians, median)
+	}
+
+	return medians
+}

runningmedian/median_test.go

@@ -0,0 +1,38 @@
+package runningmedian_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/thenativeweb/codingcircle/runningmedian"
+)
+
+func TestGetMedians(t *testing.T) {
+	t.Run("single element", func(t *testing.T) {
+		sequence := []int{23}
+		medians := runningmedian.GetMedians(sequence)
+
+		assert.Equal(t, []float64{23}, medians)
+	})
+
+	t.Run("two elements", func(t *testing.T) {
+		sequence := []int{23, 42}
+		medians := runningmedian.GetMedians(sequence)
+
+		assert.Equal(t, []float64{23, 32.5}, medians)
+	})
+
+	t.Run("multiple elements (odd)", func(t *testing.T) {
+		sequence := []int{23, 42, 36, 35, 31, 27, 172}
+		medians := runningmedian.GetMedians(sequence)
+
+		assert.Equal(t, []float64{23, 32.5, 36, 35.5, 35, 33, 35}, medians)
+	})
+
+	t.Run("multiple elements (even)", func(t *testing.T) {
+		sequence := []int{23, 42, 36, 35, 31, 27}
+		medians := runningmedian.GetMedians(sequence)
+
+		assert.Equal(t, []float64{23, 32.5, 36, 35.5, 35, 33}, medians)
+	})
+}

runningmedian/min_heap.go

@@ -0,0 +1,33 @@
+package runningmedian
+
+import "container/heap"
+
+type MinHeap []int
+
+func NewMinHeap() *MinHeap {
+	h := &MinHeap{}
+	heap.Init(h)
+	return h
+}
+
+func (h MinHeap) Len() int {
+	return len(h)
+}
+func (h MinHeap) Less(i, j int) bool {
+	return h[i] < h[j]
+}
+func (h MinHeap) Swap(i, j int) {
+	h[i], h[j] = h[j], h[i]
+}
+
+func (h *MinHeap) Push(value any) {
+	*h = append(*h, value.(int))
+}
+
+func (h *MinHeap) Pop() any {
+	oldHeap := *h
+	valueIndex := len(oldHeap) - 1
+	value := oldHeap[valueIndex]
+	*h = oldHeap[:valueIndex]
+	return value
+}