Cases versus Invariants
I stumbled upon a situation where I learnt about invariants to write less diabolical code.
When I talk about case-based-reasoning I mean approaching a problem by breaking it down into cases ie., if this happens do this, if that happens do that
You read more about invariants here they probably have a nice explanation over there. It basically means a condition or a property that must be true and always true throughout the algorithm.
Now, let’s take the situation ie., finding the median in the stream of data.
The Situation
We would start out with an empty array or a slice and we need to implement a couple of methods ie., adding a number and returning the median at that time.
So, to find the median we can either sort the array each time or we can traverse through the array, to find the right position to insert the number and move all the other elements to the right.
That brings us to either O(n log n) or O(n) for each insertion, and O(1) for returning the median.
And yeah, we can do better.
So the trick is to use heaps or priority queues.
We would maintain two heaps:
- The maxheap the smaller half.
- The minheap the bigger half.
And this is where the invariants come in:
- The heaps differ in size by at most 1.
- Every element in the minheap is less than or equal to every element in the maxheap.
So whenever we need the median we can do:
- If there are odd number of elements, the median is the root of the maxheap.
- If there are even number of elements, the median is the average of the roots of both the heaps.
The Solution
Before implementing AddNum we must implement things required for it.
Since we need heaps and I wrote it. Go has container/heap and it expects us to satisfy heap.Interface.
package main
type minheap []int
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(x any) {
(*h) = append((*h), x.(int))
}
func (h *minheap) Pop() any {
n := len(*h)
x := (*h)[n-1]
*h = (*h)[:n-1]
return x
}
type maxheap []int
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(x any) {
(*h) = append((*h), x.(int))
}
func (h *maxheap) Pop() any {
n := len(*h)
x := (*h)[n-1]
*h = (*h)[:n-1]
return x
}
Now that we have the minheap and the maxheap we can move on to implementing MedianFinder.
package main
type MedianFinder struct {
// maxheap contains the smaller half of the numbers.
// Its root is the largest number in the smaller half.
minheap *minheap
// minheap contains the bigger half of the numbers.
// Its root is therefore the smallest number in the bigger half.
maxheap *maxheap
}
// Constructor returns a MedianFinder,
// in other words a instance of medianFinder.
func Constructor() MedianFinder {
return MedianFinder{
minheap: &minheap{},
maxheap: &maxheap{},
}
}
// FindMedian returns the median of the array.
func (this *MedianFinder) FindMedian() float64 {
// when there are odd number of elements,
// maxheap has one more element than minheap
// so its root is the middle element
if n := (*this.maxheap).Len() + (*this.minheap).Len(); n%2 == 1 {
return float64((*this.maxheap)[0])
}
// when there are even number of elements,
// the roots of the both the heaps are the middle elements
// so we return their average.
l := (*this.maxheap)[0]
r := (*this.minheap)[0]
return float64(l+r) / float64(2)
}
Case-Based Reasoning
Now that we have our pre-requisites we can move on to implement AddNum.
AddNum has a simple responsibility, add a number while keeping the invariants true.
Don’t worry it’ll look diabolical, but actually it is simple.
package main
import (
"container/heap"
"math"
)
func (this *MedianFinder) AddNum(num int) {
// when both heaps have the same size,
// adding one element should make the max-heap contain one more element than the min-heap
if (*this.maxheap).Len() == (*this.minheap).Len() {
var bound int
if (*this.maxheap).Len() == 0 {
// there is no existing value to compare against
// treat the bound as +infinity
// so the first element goes directly into the max-heap
bound = math.MaxInt
} else {
// largest value in the max-heap is the boundary between
// the two halves of the numbers
bound = (*this.maxheap)[0]
}
if num < bound {
// num belongs to the lower half
// so add it to the max-heap
heap.Push(this.maxheap, num)
} else {
// num belongs to the upper half,
// so initially put it in the min-heap
//
// we then move its smallest element to the max-heap
// to keep the max-heap one element larger
//
// don't worry if the shuffling isn't clear
// there is a diagram below that explains this
heap.Push(this.minheap, num)
x := heap.Pop(this.minheap)
heap.Push(this.maxheap, x)
}
return
}
// max-heap currently contains one more element than the min-heap
// after adding the new number, both heaps should have the same size
var bound int
if (*this.maxheap).Len() == 0 {
// this case is technically unreachable here
// because the max-heap must be larger than the min-heap
bound = math.MinInt
} else {
// root of the max-heap is the largest value in the lower half.
bound = (*this.maxheap)[0]
}
if bound < num {
// num belongs to the upper half,
// so add it directly to the min-heap
heap.Push(this.minheap, num)
} else {
// num belongs to the lower half,
// so add it to the max-heap first
//
// this would make the max-heap large,
// so move its largest element to the min-heap
// to restore the size balance
//
// again, don't worry if the shuffling isn't clear
// there is a diagram below that explains this
heap.Push(this.maxheap, num)
x := heap.Pop(this.maxheap)
heap.Push(this.minheap, x)
}
}
Now the shuffling.
Sometimes an element belongs in one heap, but its value doesn’t fit the ordering of the other heap.
Instead of moving elements around manually, we temporarily put the element in the other heap.
Then pop the appropriate element and move it into the correct heap.
See the example below:
Invariant-Based Reasoning
Uhhhh, that was tiring right? so many conditions and cases and Now it is time for some invariant-based reasoning:
package main
import "container/heap"
func (this *MedianFinder) AddNum(num int) {
// first, put the element in the maxheap
//
// if the element is larger than the elements in the minheap,
// it will be moved immediately to the minheap below
heap.Push(this.maxheap, num)
// moves the largest element from the maxheap to the minheap
// this keeps every element in maxheap <= every element in minheap,
// so this takes care of the first invariant
heap.Push(this.minheap, heap.Pop(this.maxheap))
// if the minheap becomes bigger than the maxheap,
// which violates our second invariant
// simply move the smallest element from the minheap to the maxheap
if this.maxheap.Len() < this.minheap.Len() {
heap.Push(this.maxheap, heap.Pop(this.minheap))
}
}
Honestly I don’t think you need any more explanation than this.
Can you see the difference? It is literally night and day!
The End
And, sorry if I chose a relatively confusing, hard example for this. But this is where I encountered the day and night experience.
I hope you experienced that Is that it? moment too.
And finally, you definitely won’t write simple, elegant code unless you’ve written some diabolical ones first.