Search code examples
pythonlistsortingranking

Efficient method to calculate the rank vector of a list in Python, handle ties


I'm looking for an efficient way to calculate the rank vector of a list in Python, similar to R's rank function. In a simple list with no ties between the elements, element i of the rank vector of a list l should be x if and only if l[i] is the x-th element in the sorted list. This is simple so far, the following code snippet does the trick:

def rank_simple(vector):
    return sorted(range(len(vector)), key=vector.__getitem__)

Things get complicated, however, if the original list has ties (i.e. multiple elements with the same value). In that case, all the elements having the same value should have the same rank, which is the average of their ranks obtained using the naive method above. So, for instance, if I have [1, 2, 3, 3, 3, 4, 5], the naive ranking gives me [0, 1, 2, 3, 4, 5, 6], but what I would like to have is [0, 1, 3, 3, 3, 5, 6]. Which one would be the most efficient way to do this in Python?


Footnote: let me know if there's a NumPy method to achieve this; but I am interested in a pure Python solution anyway as I'm developing a tool which should work without NumPy as well.


Solution

  • Using scipy, the function you are looking for is scipy.stats.rankdata:

    In [13]: import scipy.stats as ss
    In [19]: ss.rankdata([3, 1, 4, 15, 92])
    Out[19]: array([ 2.,  1.,  3.,  4.,  5.])
    
    In [20]: ss.rankdata([1, 2, 3, 3, 3, 4, 5])
    Out[20]: array([ 1.,  2.,  4.,  4.,  4.,  6.,  7.])
    

    The ranks start at 1, rather than 0 (as in your example), but then again, that's the way R's rank function works as well.

    Here is a pure-python equivalent of scipy's rankdata function:

    def rank_simple(vector):
        return sorted(range(len(vector)), key=vector.__getitem__)
    
    def rankdata(a):
        n = len(a)
        ivec=rank_simple(a)
        svec=[a[rank] for rank in ivec]
        sumranks = 0
        dupcount = 0
        newarray = [0]*n
        for i in xrange(n):
            sumranks += i
            dupcount += 1
            if i==n-1 or svec[i] != svec[i+1]:
                averank = sumranks / float(dupcount) + 1
                for j in xrange(i-dupcount+1,i+1):
                    newarray[ivec[j]] = averank
                sumranks = 0
                dupcount = 0
        return newarray
    
    print(rankdata([3, 1, 4, 15, 92]))
    # [2.0, 1.0, 3.0, 4.0, 5.0]
    print(rankdata([1, 2, 3, 3, 3, 4, 5]))
    # [1.0, 2.0, 4.0, 4.0, 4.0, 6.0, 7.0]