NumPy: calculate cumulative median

Knowing that Python has a heapq module that lets you keep a running 'minimum' for an iterable, I did a search on heapq and median, and found various items for steaming medium. This one:

http://www.ardendertat.com/2011/11/03/programming-interview-questions-13-median-of-integer-stream/

has a class streamMedian that maintains two heapq, one with the bottom half of the values, the other with top half. The median is either the 'top' of one or the mean of values from both. The class has an insert method and a getMedian method. Most of the work is in the insert.

I copied that into an Ipython session, and defined:

def cummedian_stream(b):
    S=streamMedian()
    ret = []
    for item in b:
        S.insert(item)
        ret.append(S.getMedian())
    return np.array(ret)

Testing:

In [155]: a = np.random.randint(0,100,(5000))
In [156]: amed = cummedian_stream(a)
In [157]: np.allclose(cummedian_sorted(a), amed)
Out[157]: True
In [158]: timeit cummedian_sorted(a)
1 loop, best of 3: 781 ms per loop
In [159]: timeit cummedian_stream(a)
10 loops, best of 3: 39.6 ms per loop

The heapq stream approach is way faster.

---

The list comprehension that @Uriel gave is relatively slow. But if I substitute np.median for statistics.median it is faster than @Divakar's sorted solution:

def fastloop(a):
    return np.array([np.median(a[:i+1]) for i in range(len(a))])

In [161]: timeit fastloop(a)
1 loop, best of 3: 360 ms per loop

And @Paul Panzer's partition approach is also good, but still slow compared to the streaming class.

In [165]: timeit cummedian_partition(a)
1 loop, best of 3: 391 ms per loop

(I could copy the streamMedian class to this answer if needed).

Here's an approach that replicates elements along rows to give us a 2D array. Then, we would fill the upper triangular region with a big number so that later on when we sort the array along each row, would basically sort all elements till the diagonal elements and that simulates the cumulative windows. Then, following the definition of median that chooses the middle one or the mean of two middle ones (for even no. of elements), we would get the elements at the first position : (0,0), then for the second row : mean of (1,0) & (1,1), for the third row : (2,1), for the fourth row : mean of (3,1) & (3,2) and so on. So, we will extract out those elements from the sorted array and thus have our median values.

Thus, the implementation would be -

def cummedian_sorted(a):
    n = a.size
    maxn = a.max()+1
    a_tiled_sorted = np.tile(a,n).reshape(-1,n)
    mask = np.triu(np.ones((n,n),dtype=bool),1)

    a_tiled_sorted[mask] = maxn
    a_tiled_sorted.sort(1)

    all_rows = a_tiled_sorted[np.arange(n), np.arange(n)//2].astype(float)
    idx = np.arange(1,n,2)
    even_rows = a_tiled_sorted[idx, np.arange(1,1+(n//2))]
    all_rows[idx] += even_rows
    all_rows[1::2] /= 2.0
    return all_rows

Runtime test

Approaches -

# Loopy solution from @Uriel's soln   
def cummedian_loopy(arr):
    return [median(a[:i]) for i in range(1,len(a)+1)]

# Nan-fill based solution from @Nickil Maveli's soln   
def cummedian_nanfill(arr):
    a = np.tril(arr).astype(float)
    a[np.triu_indices(a.shape[0], k=1)] = np.nan
    return np.nanmedian(a, axis=1)

Timings -

Set #1 :

In [43]: a = np.random.randint(0,100,(100))

In [44]: print np.allclose(cummedian_loopy(a), cummedian_sorted(a))
    ...: print np.allclose(cummedian_loopy(a), cummedian_nanfill(a))
    ...: 
True
True

In [45]: %timeit cummedian_loopy(a)
    ...: %timeit cummedian_nanfill(a)
    ...: %timeit cummedian_sorted(a)
    ...: 
1000 loops, best of 3: 856 µs per loop
1000 loops, best of 3: 778 µs per loop
10000 loops, best of 3: 200 µs per loop

Set #2 :

In [46]: a = np.random.randint(0,100,(1000))

In [47]: print np.allclose(cummedian_loopy(a), cummedian_sorted(a))
    ...: print np.allclose(cummedian_loopy(a), cummedian_nanfill(a))
    ...: 
True
True

In [48]: %timeit cummedian_loopy(a)
    ...: %timeit cummedian_nanfill(a)
    ...: %timeit cummedian_sorted(a)
    ...: 
10 loops, best of 3: 118 ms per loop
10 loops, best of 3: 47.6 ms per loop
100 loops, best of 3: 18.8 ms per loop

Set #3 :

In [49]: a = np.random.randint(0,100,(5000))

In [50]: print np.allclose(cummedian_loopy(a), cummedian_sorted(a))
    ...: print np.allclose(cummedian_loopy(a), cummedian_nanfill(a))

True
True

In [54]: %timeit cummedian_loopy(a)
    ...: %timeit cummedian_nanfill(a)
    ...: %timeit cummedian_sorted(a)
    ...: 
1 loops, best of 3: 3.36 s per loop
1 loops, best of 3: 583 ms per loop
1 loops, best of 3: 521 ms per loop

Use statistics.median and cummulative list comprehension (note that odd indices contains medians of even-length lists - where the median is the average of the two median elements, so it usually results with a decimal and not an integer):

>>> from statistics import median
>>> arr = [1, 3, 4, 2, 5, 3, 6]
>>> cum_median = [median(arr[:i+1]) for i in range(len(arr)-1)]
>>> cum_median
[1, 2.0, 3, 2.5, 3, 3.0]