I am new in Python and I exercising in writing codes, but I am having some troubles.
I am trying to implement an algorithm based on QuickSelect in which it's possible to extract the K largest elements.
So, before I implement the QuickSelect
algorithm and use it to find the K-th largest element in the array A. Then using the function k_largest_quickselect(A, K)
I scan the array A to collect the K elements larger or equal to the K-th element found with QuickSelect
Finally, sort the collected elements.
This is the code of my implementation:
def partition(A, left, right):
pivot = random.choice(A) # 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)
def k_largest_quickselect(A, K):
B = [i for i in A if i >= QuickSelect(A = A, K = K, left = 0, right = len(a)-1)] # elements >= the K-th element found with QuickSelect
B.sort(reverse = True)
return B
I tried to test the algorithm
a = get_random_array(10, 10)
print("Array a = " + str(a))
print(sorted(a)[-3:])
print(sorted(k_largest_quickselect(a, 3))) # from array a select the 3 highest number
getting this result
Array a = [0, 3, 0, 6, 2, 5, 1, 8, 1, 9]
[6, 8, 9]
[2, 5, 6, 6, 9, 9]
As you can see, by using the function k_largest_quickselect(A = a, K = 3)
I did not get the 3 largest elements of the array.
Please, could you help me to solve this problem? Thank you very much!