Find the element of rank k using qsort type algorithm

kthelement_qsort.py

def find_kth_element(a, lo, hi, k):
  
    if lo >= hi:
        return a[lo]

    p = partition(a, lo, hi)

    if (k-1) == p:
        return a[p]

    if (k - 1) < p:
        return find_kth_element(a, lo, p-1, k)

    if (k + 1) > p:
        return find_kth_element(a, p+1, hi, k)


def partition(a, lo, hi):
   
    pivot = a[hi]
    i = lo
    j = lo

    while j <= hi:
        if (a[j] < pivot):
            a[i], a[j] = a[j], a[i]
            i += 1
        j += 1

    a[i], a[hi] = a[hi], a[i]

    return i