Rolling average for sorted and unsorted streams

AAA_Rolling_Average.txt

Rolling average for sorted and unsorted streams 

MaxHeapComparator.java

package com.blogspot.groglogs.rollingmedian.core;

import java.util.Comparator;

public class MaxHeapComparator implements Comparator<Float>{

    @Override
    /*
    -1 x < y
    0 x = y
    1 x > y
    For a max heap we invert this
    */
    public int compare(Float x, Float y)
    {
        if (x < y)
        {
            return 1;
        }
        if (x > y)
        {
            return -1;
        }
        return 0;
    }
}

RollingMedian.java

package com.blogspot.groglogs.rollingmedian;

import com.blogspot.groglogs.rollingmedian.core.MaxHeapComparator;

import java.util.PriorityQueue;
import java.util.Queue;
import java.util.LinkedList;

public class RollingMedian {

    //unsorted median
    private static int count = 0;
    private static Float odd = null, even = null;
    private static Queue<Float> next = new LinkedList<Float>();

    //sorted median
    private static PriorityQueue<Float> min = new PriorityQueue<>();
    private static PriorityQueue<Float> max = new PriorityQueue<Float>(new MaxHeapComparator());
    private static Float tmp = null;

    //needed to restart the count for a completely new stream
    public static void resetUnsorted(){
        count = 0;
        odd = null;
        even = null;
        next = new LinkedList<Float>();
    }

    /*
    rolling median is median element when stream length is odd: 1,2,3 -> 2
    otherwise average of median elements when stream length is even: 1,2,3,4 -> (2+3)/2
    */
    public static float getUnsortedMedian(float next_item){
        next.add(next_item);
        count++;

        //odd stream length
        if(count % 2 != 0){
            if(even == null) odd = next.remove();
            else odd = even;
            return odd;
        }
        else{
            even = next.remove();
            return (odd + even) / 2;
        }
    }

    //needed to restart the count for a completely new stream
    public static void resetSorted(){
        min = new PriorityQueue<>(); //elements greater than current median
        max = new PriorityQueue<Float>(new MaxHeapComparator()); //elements smaller than current median
        tmp = null;
    }

    /*
    sorted median requires us to track the elements in sorted order as soon as they appear
    we keep a max heap for the left side of the median (elements in increasing order BEFORE median with biggest at top)
    and a min heap for the right side (elements in increasing order AFTER median with smallest at top)
    we start by always adding in left side and balance out with right allowing the left side to be at most 1 element bigger than the right side
    and never allowing the right side to have a smaller element as root than the left side
    when the count of the elements in both heaps is the same, the median is the average of the roots
    otherwise it is the element in the max heap
    */
    public static float getSortedMedian(float next_item){
        //start by always adding left
        max.add(next_item);

        //do not allow left to grow more than 1 element than right, if so, move from left to right
        if(max.peek() != null && max.size() - min.size() >= 2){
            min.add(max.poll());
        }
        //CAUTION: otherwise move anyway if top(right) is smaller than top(left) to keep the heaps balanced!
        else if (min.peek() != null && max.peek() != null && min.peek() < max.peek()){
            tmp = min.poll();
            min.add(max.poll());
            max.add(tmp);
        }

        //if both heaps have same size, the median is the average of both tops (max BEFORE median and min AFTER median)
        if(max.size() == min.size()) return (max.peek() + min.peek()) / 2;
        //otherwise we are already tracking the median in the left heap
        else return max.peek();
    }
}

RollingMedianJTests.java

package com.blogspot.groglogs.rollingmedian.test;

import org.junit.Test;
import com.blogspot.groglogs.rollingmedian.RollingMedian;

import static org.junit.Assert.assertEquals;

public class RollingMedianJTests {

    float[] stream = null, out = null;
    float output = 0, current = 0, epsilon = 0.001f;

    @Test
    //one element stream - unsorted
    public void oneUnsortedElement() {
        stream = new float[]{1};
        output = 1;

        RollingMedian.resetUnsorted();

        assertEquals("unsorted 1 = 1", output, RollingMedian.getUnsortedMedian(stream[0]), epsilon);
    }

    @Test
    //multiple elements stream - odd length, test only final result - unsorted
    public void multipleUnsortedElementsOddFinal() {
        stream = new float[]{1, 2, 3, 4, 5};
        output = 3;

        RollingMedian.resetUnsorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getUnsortedMedian(stream[i]);
        }

        assertEquals("unsorted 1, 2, 3, 4, 5 = 3", output, current, epsilon);
    }

    @Test
    //multiple elements stream - even length, test only final result - unsorted
    public void multipleUnsortedElementsEvenFinal() {
        stream = new float[]{1, 2, 3, 4};
        output = 2.5f;

        RollingMedian.resetUnsorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getUnsortedMedian(stream[i]);
        }

        assertEquals("unsorted 1, 2, 3, 4 = 2.5", output, current, epsilon);
    }

    @Test
    //multiple elements stream - odd length, test each result - unsorted
    public void multipleUnsortedElementsOddEach() {
        stream = new float[]{1, 2, 3, 4, 5};
        out = new float[]{1, 1.5f, 2, 2.5f, 3};

        RollingMedian.resetUnsorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getUnsortedMedian(stream[i]);
            assertEquals("unsorted rolling median i = " + i, out[i], current, epsilon);
        }
    }

    @Test
    //multiple elements stream - even length, test each result - unsorted
    public void multipleUnsortedElementsEvenEach() {
        stream = new float[]{1, 2, 3, 4};
        out = new float[]{1, 1.5f, 2, 2.5f};

        RollingMedian.resetUnsorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getUnsortedMedian(stream[i]);
            assertEquals("unsorted rolling median i = " + i, out[i], current, epsilon);
        }
    }

    @Test
    //one element stream - sorted
    public void oneSortedElement() {
        stream = new float[]{1};
        output = 1;

        RollingMedian.resetSorted();

        assertEquals("sorted 1 = 1", output, RollingMedian.getSortedMedian(stream[0]), epsilon);
    }

    @Test
    //multiple elements stream - odd length, test only final result - sorted
    public void multipleSortedElementsOddFinal() {
        stream = new float[]{1, 2, 3, 4, 5};
        output = 3;

        RollingMedian.resetSorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getSortedMedian(stream[i]);
        }

        assertEquals("sorted 1, 2, 3, 4, 5 = 3", output, current, epsilon);
    }

    @Test
    //multiple elements stream - even length, test only final result - sorted
    public void multipleSortedElementsEvenFinal() {
        stream = new float[]{1, 2, 3, 4};
        output = 2.5f;

        RollingMedian.resetSorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getSortedMedian(stream[i]);
        }

        assertEquals("sorted 1, 2, 3, 4 = 2.5", output, current, epsilon);
    }

    @Test
    //multiple elements stream - odd length, test each result - sorted
    public void multipleSortedElementsOddEach() {

        stream = new float[]{9, 4, 10, 3, 3, 14, 14, 0, 16};
        out = new float[]{9, 6.5f, 9, 6.5f, 4, 6.5f, 9, 6.5f, 9};

        RollingMedian.resetSorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getSortedMedian(stream[i]);
            assertEquals("sorted rolling median i = " + i, out[i], current, epsilon);
        }
    }

    @Test
    //multiple elements stream - even length, test each result - sorted
    public void multipleSortedElementsEvenEach() {

        stream = new float[]{9, 4, 10, 3, 3, 14, 14, 0, 16, 5};
        out = new float[]{9, 6.5f, 9, 6.5f, 4, 6.5f, 9, 6.5f, 9, 7};

        RollingMedian.resetSorted();

        for(int i = 0; i < stream.length; i++){
            current = RollingMedian.getSortedMedian(stream[i]);
            assertEquals("sorted rolling median i = " + i, out[i], current, epsilon);
        }
    }
}

Full description at: http://groglogs.blogspot.ch/2017/04/java-rolling-average-for-both-sorted.html