选择第n小的元素之python实现源码

def partition(A, p, r):
    j = p+1
    for i in range(p+1, r+1):
        if(A[i] < A[p]):
            tmp = A[i]
            A[i] = A[j]
            A[j] = tmp
            j = j + 1

    tmp = A[p]
    A[p] = A[j-1]
    A[j-1] = tmp
    return j - 1
   


def select(A, p, r, i):
    if(p == r):
        return A[p]
    q = partition(A, p, r)
    #print q
    k = q - p + 1
    if(k == i):
        return A[q]
    if(i < k):
        return select(A, p, q - 1, i)
    else:
        return select(A, q + 1, r, i - k)

#test case
#import pdb
#pdb.set_trace()
old = [2, 5, 3, 0, 7, 9, 6, 3]     
q = select(old, 0, 7, 6)
print q