Find K-th largest element with QuickSelect in Python

I am new in Python and I exercising in writing codes, but I am having some troubles.

I am trying to implement QuickSelect in order to extract the K largest element.

This is my code;

def partition(A, left, right): 
    pivot = random.randrange(left, right)  # pick a random number as pivot
    i = left - 1
    for j in range(left, right): 
        if A[j] <= pivot: 
            i = i+1 
            A[i], A[j] = A[j], A[i]
    A[i+1], A[right] = A[right], A[i+1] 
    return i+1

def QuickSelect(A, K, left, right): # Array, K-th element
    if left == right:
        return A[left]
    q = partition(A, left, right)
    i = q - left + 1
    if K == i:
        return A[i]
    if K < i:
        return QuickSelect(A, K, left, q - 1)
    else:
        return QuickSelect(A, K - i, q + 1, right)

Trying to implement the algorithm in order to get the 5-th highest element :

a = get_random_array(10, 100)
print("array unsorted=" ,  a)
print("array sorted=", sorted(a))
print("result:" , QuickSelect(A = a, K = 5, left = 0, right = len(a)-1)) #I want the 5-th highest element

getting this result:

array unsorted = [71, 34, 0, 36, 26, 15, 3, 69, 93, 85]
array sorted = [0, 3, 15, 26, 34, 36, 69, 71, 85, 93]
result: 3

The result is wrong since 3 is not the 5-th highest element.

Is the error in partition(A, left, right) or in QuickSelect(A, K, left, right)? Could you help me to solve it please? Thank you!

Solution

There are several issues in your code.

The code treats pivot as a value, but it is an index.
The pivot should first be moved out of the way (to right), before you start the loop
When K == i you should not return A[i], as i is a relative, 1-based index. You should return A[q]
Not a serious problem, but randrange(left, right) will never return right, so probably you want to do randrange(left, right + 1) or randint(left, right)

The corrected code:

def partition(A, left, right): 
    # This gets an index, not a value:
    pivotindex = random.randint(left, right)  # allow right to be selected
    pivot = A[pivotindex]  # this is the pivot value
    # move the value out of the way
    A[right], A[pivotindex] = A[pivotindex], A[right]
    i = left - 1
    for j in range(left, right): 
        if A[j] < pivot:
            i += 1 
            A[i], A[j] = A[j], A[i]
    A[i+1], A[right] = A[right], A[i+1] 
    return i + 1

def QuickSelect(A, K, left, right):
    if left == right:
        return A[left]
    q = partition(A, left, right)
    i = q - left + 1
    if K == i:
        return A[q]  # this is the element you want to return
    if K < i:
        return QuickSelect(A, K, left, q - 1)
    else:
        return QuickSelect(A, K - i, q + 1, right)