RSelect.py

from random import randint
import random


def get_pivot(low, max):
    return randint(low, max)


def partition(ar, low, max, p):
    ar[low], ar[p] = ar[p], ar[low]
    p = ar[low]
    i = low + 1
    for j in range(i, max + 1):
        if ar[j] < p:
            ar[i], ar[j] = ar[j], ar[i]
            i += 1
    ar[i - 1], ar[low] = ar[low], ar[i - 1]
    return i - 1


def rselect(ar, low, max, i):
    if low >= max:
        return ar[low]
    p = get_pivot(low, max)
    j = partition(ar, low, max, p)
    
    if j == i:
        return ar[j]
    elif j > i:
        return rselect(ar, low, j - 1, i)
    else:
        return rselect(ar, j + 1, max, i)


ar = [5, 4, 9, 10, 1, 6]

istat = 3
calculatedIstat = rselect(ar, 0, len(ar) - 1, istat)
print("{}-ith in {} is {}".format(istat, ar, calculatedIstat))
assert calculatedIstat == ar[istat], "calculatedIstat(value={}) should be ar[{}]={}".format(calculatedIstat, istat - 1, ar[istat])



for i in range(1000):
    randnums = random.sample(range(1000), 700)
    #sorted
    sortedRandNums = sorted(randnums)
    selectedSortedRandNum = sortedRandNums[44]

    #my select
    selectedRandNum = rselect(randnums, 0, len(randnums) - 1, 44)

    assert selectedRandNum == selectedSortedRandNum