Find K-th largest element with QuickSelect in Python

Question

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!

This code doesn't run: you pass `i = 5` as argument to `QuickSelect`, but it doesn't have this named argument. — trincot, Jun 13 '21 at 12:35
@trincot I fixed the problem of ```i = 5``` by passing ```K = 5```, but still the same problem. Is there any error with indexes in the functions? — user15874067, Jun 13 '21 at 12:41

Varun Singh · Answer 1 · 2021-06-13T12:52:42.557

You are having problem with both of your functions. You are setting pivot as in partition function

random.randrange(left, right)

instead set it as

A[right]

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

In the second function re-modifiy it as below

def Quickselect(A, left, right, k): 

    if (k > 0 and k <= right - left + 1):
  
        q = partition(A, left, right)
        if (q - left == k - 1):
            return A[q] 
        if (q - left > k - 1):
            return Quickselect(A, left, q - 1, k) 
        return Quickselect(A, q + 1, right,
                            k - q + left - 1)

Reference: https://www.geeksforgeeks.org/quickselect-algorithm/

Blockquote

score 0 · Accepted Answer · answered Jun 13 '21 at 12:58

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)

Find K-th largest element with QuickSelect in Python

2 Answers2