Day 18 - Selection algorithm


Definition(s)

Selection algorithm is an algorithm for finding the kth smallest number in a list or array; such a number is called the kth order statistic.

The best-known selection algorithm is quickselect, which is related to quicksort;

Algorithm(s)


In [1]:
import numpy as np

In [2]:
def impl_nth_element(A, l, r, k):
    if l < r:
        i, j, = l, r
        pivot = A[np.random.randint(l, r, size=1)]
        
        while True:
            while A[i] < pivot:
                i += 1
                
            while A[j] > pivot:
                j -= 1
                
            if i <= j:
                A[i], A[j] = A[j], A[i]
                i += 1
                j -= 1
                
            if i > j:
                break
                
        if l <= k <= j:
            impl_nth_element(A, l, j, k)
            
        if i <= k <= r:
            impl_nth_element(A, i, r, k)

In [3]:
def nth_element(a, k):
    impl_nth_element(a, 0, len(a) - 1, k)
    return a[k]

def median(a):
    return nth_element(a, len(a) // 2)

In [4]:
import heapq

def heap_select(a, k):
    return heapq.nsmallest(k + 1, a)[-1]

Run(s)


In [5]:
np.random.seed(0xDEAD)

In [10]:
xs = np.random.randint(0, 10, size=10)
print(xs)
print(nth_element(xs, 2))
print(heap_select(xs, 2))
print(xs)


[1 1 4 2 0 9 9 0 3 9]
1
1
[0 0 1 1 2 9 9 4 3 9]

In [7]:
xs = np.random.randint(0, 100, size=20)
print(xs)
print(nth_element(xs, 4))
print(heap_select(xs, 4))
print(xs)


[59 15 71 86 64 57 23 72 29 45 66 47 64 39 31 92 53 90 25 34]
31
31
[29 15 25 23 31 39 53 47 34 45 66 72 64 57 64 92 86 90 71 59]

In [13]:
xs = np.random.randint(0, 100, size=25)
print(xs)
print(nth_element(xs, 15))
print(heap_select(xs, 15))
print(xs)


[69  7 99 42 96 84 13 38 86 71 19 93 43 51 11 72 18 26 27 99 88 13 65 38 21]
69
69
[26  7 21 18 11 13 13 19 65 27 38 42 43 51 38 69 72 84 71 99 88 93 86 96 99]

In [9]:
xs = np.array(range(10))
print(xs)
print(nth_element(xs, 2))
print(heap_select(xs, 2))
print(xs)


[0 1 2 3 4 5 6 7 8 9]
2
2
[0 1 2 3 4 5 6 7 8 9]